feat(server): add --api-key bearer authentication

docs/cli.md

@@ -40,6 +40,7 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--host` | 127.0.0.1 | Bind address |
 | `--port` | 1919 | Bind port |
+| `--api-key` | disabled | Require `Authorization: Bearer <key>` on every route except `/health` (401 otherwise); `FREETOKEN_API_KEY` is read when the flag is absent |
 | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) |
 | `--max-running-requests` | 4 | Max concurrently running requests |
 | `--max-output-tokens` | 32768 | Default output budget for requests that omit one |

python/freetoken/server/api_server.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import contextlib
+import hmac
 import json
 import os
 import signal
@@ -437,6 +438,23 @@ def install_cors(app: FastAPI, origins_csv: str) -> None:
 _UNTRACKED_REQUEST_PREFIXES = ("/v1/messages/count_tokens",)
 
 
+# --api-key (ServerArgs.api_key), installed by run_api_server before serving. None means no
+# authentication -- every route answers as before. Module-level rather than closed over so the
+# middleware below can be registered at import time like the others and tests can set it.
+_API_KEY: str | None = None
+# Routes that stay open with a key set: liveness probes (load balancers, Docker healthchecks,
+# the desktop app's load-progress polling) cannot carry a header and reveal nothing but status.
+_API_KEY_OPEN_PATHS = ("/health",)
+
+
+def install_api_key(key: str | None) -> None:
+    """Arm the bearer check for every route but ``_API_KEY_OPEN_PATHS``; ``None``/"" disarms."""
+    global _API_KEY
+    _API_KEY = key or None
+    if _API_KEY is not None:
+        logger.info("API key required (Authorization: Bearer) on every route except /health")
+
+
 def _served_model_name() -> str | None:
     st = _GLOBAL_STATE
     cfg = getattr(st, "config", None) if st is not None else None
@@ -566,6 +584,36 @@ def _resolve_num_swa_pages(state: FrontendManager, req: CacheRebuildRequest) ->
     return max(1, -(-window_tokens // swa_page_size))  # ceil-div to the pool's page unit
 
 
+@app.middleware("http")
+async def _api_key_middleware(request: Request, call_next):
+    """Reject any request without ``Authorization: Bearer <api_key>`` when a key is set.
+
+    Registered after ``_record_request_middleware`` so it runs *before* it (Starlette wraps the
+    last-added middleware outermost): a 401 never lands in the request ring or a handler. CORS
+    preflights (OPTIONS) carry no credentials by design and pass through; the CORS middleware
+    installed at startup is outer still, so it answers them. Constant-time compare, and the
+    same body shape the OpenAI-compatible routes use for errors."""
+    key = _API_KEY
+    if key is None or request.method == "OPTIONS" or request.url.path in _API_KEY_OPEN_PATHS:
+        return await call_next(request)
+    scheme, _, token = request.headers.get("authorization", "").partition(" ")
+    if scheme.lower() != "bearer" or not hmac.compare_digest(
+        token.strip().encode("utf-8"), key.encode("utf-8")
+    ):
+        return JSONResponse(
+            status_code=401,
+            headers={"WWW-Authenticate": "Bearer"},
+            content={
+                "error": {
+                    "message": "Invalid or missing API key (Authorization: Bearer <key>).",
+                    "type": "authentication_error",
+                    "code": 401,
+                }
+            },
+        )
+    return await call_next(request)
+
+
 @app.post("/v1/cache/rebuild")
 async def cache_rebuild(req: CacheRebuildRequest):
     """Trigger a runtime KV/MoE cache resize. Blocks until the scheduler reports a result
@@ -873,7 +921,7 @@ def _flag_shutdown(signum, frame) -> None:
         signal.signal(sig, _flag_shutdown)
 
 
-def _serve_and_run_shell(host: str, port: int) -> None:
+def _serve_and_run_shell(host: str, port: int, api_key: str | None = None) -> None:
     """Shell mode: serve the API here, and attach the terminal client to it over the loopback.
 
     The shell is an ordinary API client (see ``freetoken.shell``), so shell mode is just
@@ -906,7 +954,8 @@ def _serve_and_run_shell(host: str, port: int) -> None:
         # /health and echoes the same load progress the desktop app polls for. A ^C during that
         # wait is a stop, not a crash -- exit through the teardown below, not a traceback.
         with contextlib.suppress(KeyboardInterrupt):
-            asyncio.run(run_shell(origin, connect_grace=30.0))
+            # The attached shell is an ordinary client, so it carries the server's own key.
+            asyncio.run(run_shell(origin, connect_grace=30.0, api_key=api_key))
     finally:
         server.should_exit = True
         thread.join(timeout=15)
@@ -951,6 +1000,7 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s
     # 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.
     install_cors(app, config.cors_origins)
+    install_api_key(config.api_key)
     init_request_logging()
     # Hide the frequent health/stats/requests/cache-status polling of the desktop app (and of
     # the shell's status bar) from uvicorn's access log; non-polling access lines are
@@ -1034,7 +1084,7 @@ def _on_meta(meta: dict) -> None:
     ).start()
 
     if run_shell:
-        _serve_and_run_shell(host, port)
+        _serve_and_run_shell(host, port, config.api_key)
         return
     # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group.
     uvicorn.run(app, host=host, port=port)

python/freetoken/server/args.py

@@ -15,6 +15,9 @@
 class ServerArgs(SchedulerConfig):
     server_host: str = "127.0.0.1"
     server_port: int = 1919
+    # Bearer token every request must carry (except /health). None = no authentication,
+    # today's behaviour. Read from FREETOKEN_API_KEY when --api-key is not given.
+    api_key: str | None = None
     num_tokenizer: int = 0
     silent_output: bool = False
     # The terminal shell is attached to this server (ft shell --model / ft serve --shell-mode).
@@ -308,6 +311,17 @@ 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=ServerArgs.api_key,
+        help=(
+            "Require `Authorization: Bearer <key>` on every route except /health "
+            "(401 otherwise). Unset: no authentication. When the flag is absent, "
+            "FREETOKEN_API_KEY is read instead so the key need not appear in `ps`."
+        ),
+    )
+
     parser.add_argument(
         "--cuda-graph-max-bs",
         "--graph",
@@ -654,6 +668,11 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
     if kwargs["model_path"].startswith("~"):
         kwargs["model_path"] = os.path.expanduser(kwargs["model_path"])
 
+    if kwargs["api_key"] is None:
+        kwargs["api_key"] = os.environ.get("FREETOKEN_API_KEY") or None
+    elif not kwargs["api_key"].strip():
+        parser.error("--api-key must not be empty (omit it to serve without authentication)")
+
     if kwargs["served_model_name"] is None:
         kwargs["served_model_name"] = (
             os.path.basename(os.path.normpath(kwargs["model_path"])) or kwargs["model_path"]

python/freetoken/shell/client.py

@@ -102,6 +102,10 @@ def __init__(
     ) -> None:
         self.origin = origin.rstrip("/")
         self.timeout = timeout
+        # Sent on the control plane too (the OpenAI client already sends it on chat), so a
+        # server started with --api-key accepts both halves of the shell. The default is the
+        # placeholder an unauthenticated server ignores.
+        self.api_key = api_key
         self._openai = AsyncOpenAI(
             base_url=f"{self.origin}/v1",
             api_key=api_key,
@@ -118,7 +122,7 @@ def _request_json_blocking(
         self, method: str, path: str, body: dict[str, Any] | None, timeout: float
     ) -> dict[str, Any]:
         data = None
-        headers = {"Accept": "application/json"}
+        headers = {"Accept": "application/json", "Authorization": f"Bearer {self.api_key}"}
         if body is not None:
             data = json.dumps(body).encode("utf-8")
             headers["Content-Type"] = "application/json"

python/freetoken/shell/tui.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import asyncio
+import os
 import contextlib
 import re
 import shutil
@@ -33,6 +34,7 @@
 from prompt_toolkit.styles import Style
 
 from .client import (
+    LOCAL_API_KEY,
     ContentDelta,
     ReasoningDelta,
     Sampling,
@@ -515,13 +517,19 @@ def _format_load_progress(doc: dict) -> str:
     return f"loading ({phase})..."
 
 
-async def run_shell(origin: str, *, connect_grace: float = 0.0) -> int:
+async def run_shell(
+    origin: str, *, connect_grace: float = 0.0, api_key: str | None = None
+) -> int:
     """Attach to the FreeToken server at ``origin`` and run the terminal chat.
 
     ``connect_grace`` is how long to keep retrying a refused connection before giving up --
     left at 0 when attaching to a server the user says is already running, raised when the
-    caller just started one in this process (see ``server/api_server.py``)."""
-    client = ShellClient(origin)
+    caller just started one in this process (see ``server/api_server.py``).
+
+    ``api_key`` is the server's ``--api-key`` when it has one: passed in by shell mode, read
+    from ``FREETOKEN_API_KEY`` when attaching to a running server, else the local placeholder."""
+    key = api_key or os.environ.get("FREETOKEN_API_KEY") or LOCAL_API_KEY
+    client = ShellClient(origin, api_key=key)
     try:
         return await _run_shell(client, origin, connect_grace=connect_grace)
     finally: