feat(quota): add server‑side per‑client request quotas (requires auth)

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
per-client 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 are fully opt-in and require authentication to be
enabled. The 'sqlite' is the only supported quota `type` at this time,
any other `type` will  be rejected.

Highlights:

- Adds `QuotaMiddleware` to enforce per-client request quotas:
  - Uses `Authorization: Bearer <client_id>` (from
    AuthenticationMiddleware)
  - Tracks usage via a SQLite-based KV store
  - Returns 429 when the quota is exceeded

- Extends `ServerConfig` with a `quota` section (type + config)

- Enforces strict coupling: quotas require authentication or the server
  will fail to start

Behavior changes:
- Quotas are disabled by default unless explicitly configured
- SQLite defaults to `./quotas.db` if no DB path is set
- The server requires authentication when quotas are enabled

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:
      db_path: ./quotas.db
      quota_requests_per_day: 1000
      quota_window_seconds: 86400
```

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

Author: liangwen12year

llama_stack/distribution/datatypes.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 from enum import Enum
-from typing import Annotated, Any
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, Field
 
@@ -234,6 +234,17 @@ class AuthenticationConfig(BaseModel):
     )
 
 
+class QuotaSqliteConfig(BaseModel):
+    db_path: str = Field(default="./quotas.db", description="Path to the SQLite DB file")
+    quota_requests_per_day: int = Field(default=1000, description="Maximum requests per client per window")
+    quota_window_seconds: int = Field(default=86400, description="Quota window length in seconds")
+
+
+class QuotaConfig(BaseModel):
+    type: Literal["sqlite"] = Field(description="Quota backend type: must be 'sqlite'")
+    config: QuotaSqliteConfig
+
+
 class ServerConfig(BaseModel):
     port: int = Field(
         default=8321,
@@ -253,6 +264,10 @@ class ServerConfig(BaseModel):
         default=None,
         description="Authentication configuration for the server",
     )
+    quota: QuotaConfig | None = Field(
+        default=None,
+        description=("Per client request quota configuration. If unset or null, quotas are disabled."),
+    )
 
 
 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,93 @@
+# 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
+from datetime import datetime, 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 enforcing per client daily request quotas.
+
+    Expects Authorization: Bearer <client_id> header.
+    Tracks counts in a KV store (SQLite); returns HTTP 429 when limit is exceeded.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        kv_config: KVStoreConfig | None = None,
+        default_requests_per_day: int = 1000,
+        window_seconds: int = 86400,
+    ):
+        self.app = app
+        # if no config passed, default to on disk SQLite
+        self._kv_config = kv_config or SqliteKVStoreConfig(db_path="./quotas.db")
+        self._kv: KVStore | None = None
+        self.default_limit = default_requests_per_day
+        self.window = 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":
+            client_id = scope.get("authenticated_client_id")
+            if not client_id:
+                logger.error(
+                    "QuotaMiddleware requires an authenticated client_id but none was found in the scope. "
+                    "This likely means AuthenticationMiddleware is not installed or failed."
+                )
+                return await self._send_error(
+                    send, 500, "Quota system misconfigured: missing authenticated client identity"
+                )
+
+            key = f"quota:{client_id}:{datetime.now(timezone.utc).date().isoformat()}"
+
+            try:
+                kv = await self._get_kv()
+                prev = await kv.get(key) or "0"
+                count = int(prev) + 1
+                await kv.set(key, str(count))
+                # Note: TTL/expire is only supported on backends that implement it;
+                # for SQLite we ignore expire.
+            except Exception:
+                logger.exception("Error accessing KV store for quota")
+                return await self._send_error(send, 500, "Quota service error")
+
+            if count > self.default_limit:
+                logger.warning(
+                    "Quota exceeded for client %s: %d/%d",
+                    client_id,
+                    count,
+                    self.default_limit,
+                )
+                return await self._send_error(send, 429, "Quota exceeded")
+
+        # Pass through to downstream application
+        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

@@ -49,6 +49,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,
@@ -58,6 +59,7 @@
 
 from .auth import AuthenticationMiddleware
 from .endpoints import get_all_api_endpoints
+from .quota import QuotaMiddleware
 
 REPO_ROOT = Path(__file__).parent.parent.parent.parent
 
@@ -421,6 +423,28 @@ 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:
+            logger.error(
+                "Quota enforcement requires authentication to be enabled, but no auth config is present. "
+                "Disable quotas or configure authentication."
+            )
+            raise RuntimeError("Quota middleware requires authentication middleware to be active.")
+
+    # Enforce per-client quota (only if configured and require authentication)
+    if config.server.quota:
+        logger.info("Enabling per-client quota middleware")
+
+        quota_conf = config.server.quota.config
+
+        kv_config = SqliteKVStoreConfig(db_path=quota_conf.db_path)
+
+        app.add_middleware(
+            QuotaMiddleware,
+            kv_config=kv_config,
+            default_requests_per_day=quota_conf.quota_requests_per_day,
+            window_seconds=quota_conf.quota_window_seconds,
+        )
 
     try:
         impls = asyncio.run(construct_stack(config))

tests/unit/server/test_quota.py

@@ -0,0 +1,114 @@
+# 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 app(request):
+    """
+    Create a FastAPI app with both InjectClientIDMiddleware and QuotaMiddleware.
+    Each test gets a unique client_id for safety.
+    """
+    inner_app = FastAPI()
+
+    @inner_app.get("/test")
+    async def test_endpoint():
+        return {"message": "ok"}
+
+    # Use the test name to create a unique client_id per test
+    client_id = f"client_{request.node.name}"
+
+    app = InjectClientIDMiddleware(
+        QuotaMiddleware(
+            inner_app,
+            kv_config=SqliteKVStoreConfig(db_path="./quotas_test.db"),
+            default_requests_per_day=2,
+            window_seconds=60,
+        ),
+        client_id=client_id,
+    )
+
+    return app
+
+
+def test_quota_allows_up_to_limit(app):
+    client = TestClient(app)
+
+    resp1 = client.get("/test")
+    assert resp1.status_code == 200
+    assert resp1.json() == {"message": "ok"}
+
+    resp2 = client.get("/test")
+    assert resp2.status_code == 200
+    assert resp2.json() == {"message": "ok"}
+
+
+def test_quota_blocks_after_limit(app):
+    client = TestClient(app)
+
+    # Exceed limit: 3rd request should be throttled
+    client.get("/test")
+    client.get("/test")
+    resp3 = client.get("/test")
+    assert resp3.status_code == 429
+    assert resp3.json()["error"]["message"] == "Quota exceeded"
+
+
+def test_missing_authenticated_client_id_returns_500():
+    """
+    Confirm 500 error when QuotaMiddleware runs without authenticated_client_id.
+    """
+    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"),
+        default_requests_per_day=2,
+        window_seconds=60,
+    )
+
+    client = TestClient(app)
+
+    resp = client.get("/test")
+    assert resp.status_code == 500
+    assert "Quota system misconfigured" in resp.json()["error"]["message"]