Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 34 additions & 1 deletion python/freetoken/server/api_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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-
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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.
Expand Down
11 changes: 11 additions & 0 deletions python/freetoken/server/args.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down