From 8054a6382e7e7d3b220f4d8a2e4249afccc5ec40 Mon Sep 17 00:00:00 2001
From: santanu20 <santanu20@users.noreply.github.com>
Date: Mon, 28 Sep 2026 02:26:53 +0530
Subject: [PATCH] feat(server): --api-key flag for bearer token auth on the
 HTTP API

Adds an optional --api-key (default: FREETOKEN_API_KEY env) that gates every
non-health route behind Authorization: Bearer <key>. Unset keeps today's
no-auth behavior; /health stays open either way so supervisors can poll
readiness. The gate runs before request accounting, so rejected probes never
reach the request ring.

Implements #557.
---
 python/freetoken/server/api_server.py | 35 ++++++++++++++++++++++++++-
 python/freetoken/server/args.py       | 11 +++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py
index 5592a2b0d5..2423fdb37f 100644
--- a/python/freetoken/server/api_server.py
+++ b/python/freetoken/server/api_server.py
@@ -446,6 +446,11 @@ def _served_model_name() -> str | None:
     return getattr(cfg, "served_model_name", None)
 
 
+# Bearer token gating the HTTP API; None = auth off (upstream default). Set from
+# ServerArgs.api_key in run_api_server before uvicorn takes over.
+_API_KEY: str | None = None
+
+
 @app.middleware("http")
 async def _record_request_middleware(request: Request, call_next):
     """Time every generation request into the ring for /v1/requests + /v1/stats p95. Single-
@@ -480,6 +485,33 @@ async def _record_request_middleware(request: Request, call_next):
     return response
 
 
+@app.middleware("http")
+async def _api_key_middleware(request: Request, call_next):
+    """Reject non-health requests unless they carry Authorization: Bearer <api_key> when a
+    key is configured. Registered after _record_request_middleware, so LIFO order runs this
+    gate first and unauthorized probes never reach the request ring. /health stays open so
+    orchestrators can poll readiness without the token."""
+    if _API_KEY is None or request.url.path == "/health":
+        return await call_next(request)
+    import hmac
+
+    from fastapi.responses import JSONResponse
+
+    scheme, _, token = request.headers.get("authorization", "").partition(" ")
+    if scheme.lower() != "bearer" or not hmac.compare_digest(token, _API_KEY):
+        return JSONResponse(
+            status_code=401,
+            content={
+                "error": {
+                    "message": "Invalid or missing API key. Pass Authorization: Bearer <key>.",
+                    "type": "invalid_request_error",
+                    "code": "invalid_api_key",
+                }
+            },
+        )
+    return await call_next(request)
+
+
 class CacheRebuildRequest(BaseModel):
     moe_cache_size: int | None = None
     num_pages: int | None = None
@@ -931,7 +963,7 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s
         run_shell: If True, also attach the interactive terminal shell to the served API.
     """
 
-    global _GLOBAL_STATE, _MODEL_SAMPLING
+    global _GLOBAL_STATE, _MODEL_SAMPLING, _API_KEY
 
     if config.sampling_defaults == "model" and not config.use_dummy_weight:
         _MODEL_SAMPLING = load_generation_sampling(config.model_path)
@@ -950,6 +982,7 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s
 
     host = config.server_host
     port = config.server_port
+    _API_KEY = config.api_key
 
     # Create/validate FREETOKEN_API_LOG_DIR and start the writer thread up front, so a
     # bad path is reported at boot rather than silently on the first request.
diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py
index ab2fb9b748..1bcb5c04e5 100644
--- a/python/freetoken/server/args.py
+++ b/python/freetoken/server/args.py
@@ -52,6 +52,9 @@ class ServerArgs(SchedulerConfig):
     # "model": fill unspecified request sampling params from generation_config.json
     # (temperature/top_k/top_p), like sglang. "none": use framework defaults only.
     sampling_defaults: str = "model"
+    # Bearer token required on the HTTP API (OpenAI/Anthropic/generate routes). None
+    # disables auth entirely (today's behavior); /health stays open either way.
+    api_key: str | None = None
     # Default max output (decode) tokens for a request that omits one. None falls back to the
     # adapter's built-in default (32k).
     max_output_tokens: int | None = None
@@ -354,6 +357,14 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         help="The port number for the server to listen on.",
     )
 
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=os.environ.get("FREETOKEN_API_KEY", ServerArgs.api_key),
+        help="Bearer token required on the HTTP API; unset disables auth (previous "
+        "behavior). Defaults to the FREETOKEN_API_KEY environment variable.",
+    )
+
     parser.add_argument(
         "--cuda-graph-max-bs",
         "--graph",