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 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 .", + "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",