From 8a574de915a409e629af10922ac4b3c40c816ad5 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Thu, 9 Apr 2026 20:24:19 +0000 Subject: [PATCH] feat: persistent Access Endpoint API key with Regenerate button The API key used to live only in process memory, so restarting Studio, switching models, or toggling Disable/Enable silently invalidated every external client that was using it. Users had to reopen the dialog, copy the new key, and update their scripts each time. Key state is now stored at ~/.unsloth/studio/auth/access_endpoint.json (atomic write, mode 0600). The file survives backend restarts and model load/unload. Disable preserves the key on disk and only clears it from app.state, so Enable brings the same key back. Loading a different GGUF keeps the endpoint running transparently under the new model identifier. Rotation is an explicit user action: a new JWT-protected POST /access-endpoint/regenerate route mints a fresh key and is wired to a "Regenerate" button next to the API Key field in the dialog. Deliberately not dual-auth so a leaked key cannot rotate itself. --- studio/backend/main.py | 21 ++++ studio/backend/routes/inference.py | 100 +++++++++++++++--- studio/backend/utils/paths/storage_roots.py | 4 + .../frontend/src/features/chat/chat-page.tsx | 46 ++++++-- 4 files changed, 151 insertions(+), 20 deletions(-) diff --git a/studio/backend/main.py b/studio/backend/main.py index ad19ee9679..a14e8931ff 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -131,6 +131,27 @@ async def lifespan(app: FastAPI): print("=" * 60 + "\n") else: app.state.bootstrap_password = storage.get_bootstrap_password() + + # Restore access endpoint API key from disk (if one was left enabled). + # Persisted at ~/.unsloth/studio/auth/access_endpoint.json so the key + # survives backend restarts and external clients don't need to re-fetch + # it every time Studio is restarted. + try: + from routes.inference import _read_endpoint_state + + _ep_state = _read_endpoint_state() + if _ep_state and _ep_state["enabled"]: + app.state.external_api_key = _ep_state["api_key"] + else: + app.state.external_api_key = None + except Exception as exc: + import structlog + + structlog.get_logger(__name__).warning( + "Failed to restore access endpoint state: %s", exc + ) + app.state.external_api_key = None + yield # Cleanup _hw_module.DEVICE = None diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a5957f414e..8968477309 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -303,9 +303,6 @@ async def load_model( inference_config = load_inference_config(config.identifier) - # Clear any previous API key — user must explicitly enable endpoint - fastapi_request.app.state.external_api_key = None - return LoadResponse( status = "loaded", model = config.identifier, @@ -444,9 +441,6 @@ async def load_model( except Exception: pass - # Clear any previous API key — user must explicitly enable endpoint - fastapi_request.app.state.external_api_key = None - return LoadResponse( status = "loaded", model = config.identifier, @@ -539,9 +533,6 @@ async def unload_model( Routes to the correct backend (llama-server for GGUF, Unsloth otherwise). """ try: - # Clear API key for external access - fastapi_request.app.state.external_api_key = None - # Check if the GGUF backend has this model loaded or is loading it llama_backend = get_llama_cpp_backend() if llama_backend.is_active and ( @@ -926,6 +917,41 @@ def _get_loaded_model_name(): return "default" +def _read_endpoint_state() -> dict | None: + """Return {'api_key': str, 'enabled': bool} or None if file is missing/corrupt.""" + from utils.paths.storage_roots import access_endpoint_state_path + + path = access_endpoint_state_path() + if not path.exists(): + return None + try: + data = json.loads(path.read_text()) + if isinstance(data, dict) and isinstance(data.get("api_key"), str): + return { + "api_key": data["api_key"], + "enabled": bool(data.get("enabled", False)), + } + except Exception as exc: + logger.warning("Failed to read access_endpoint.json: %s", exc) + return None + + +def _write_endpoint_state(api_key: str, enabled: bool) -> None: + """Atomically persist the endpoint state with 0600 perms.""" + from utils.paths.storage_roots import access_endpoint_state_path, ensure_dir + + path = access_endpoint_state_path() + ensure_dir(path.parent) + tmp = path.with_suffix(".json.tmp") + tmp.write_text(json.dumps({"api_key": api_key, "enabled": enabled})) + os.chmod(tmp, 0o600) + os.replace(tmp, path) # atomic on POSIX + + +def _mint_api_key() -> str: + return f"sk-unsloth-{_secrets.token_urlsafe(32)}" + + @router.get("/access-endpoint") async def get_access_endpoint( request: Request, @@ -949,14 +975,20 @@ async def enable_access_endpoint( request: Request, current_subject: str = Depends(get_current_subject), ): - """Enable external API access: generate an API key for the running llama-server.""" + """Enable external API access. + + Reuses the persisted API key if one already exists on disk; only mints a + new one on the very first enable or after the file has been removed. + llama-server is already running with --parallel 4, so concurrent Studio + chat + external API is supported without a restart. + """ llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: raise HTTPException(status_code = 400, detail = "No GGUF model loaded") - # llama-server is already running with --parallel 4, so concurrent Studio - # chat + external API is supported without a restart. We just mint a key. - api_key = f"sk-unsloth-{_secrets.token_urlsafe(32)}" + state = _read_endpoint_state() + api_key = state["api_key"] if state else _mint_api_key() + _write_endpoint_state(api_key, enabled = True) request.app.state.external_api_key = api_key local_url, external_url = _build_endpoint_urls(request) @@ -974,11 +1006,51 @@ async def disable_access_endpoint( request: Request, current_subject: str = Depends(get_current_subject), ): - """Disable external API access: clear the API key. llama-server keeps running with --parallel 4.""" + """Disable external API access. + + Preserves the key on disk (so the next Enable brings the same key back) + but clears it from in-memory app.state, which causes the dual-auth check + to reject the key on subsequent requests. + """ + state = _read_endpoint_state() + if state: + _write_endpoint_state(state["api_key"], enabled = False) request.app.state.external_api_key = None return {"enabled": False} +@router.post("/access-endpoint/regenerate") +async def regenerate_access_endpoint( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """Mint a fresh API key, invalidating the previous one. + + JWT-only (deliberately not dual-auth) so a leaked external API key + cannot rotate itself. Only legal while the endpoint is currently enabled + and a GGUF model is loaded. + """ + if getattr(request.app.state, "external_api_key", None) is None: + raise HTTPException(status_code = 400, detail = "Endpoint is not enabled") + + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException(status_code = 400, detail = "No GGUF model loaded") + + api_key = _mint_api_key() + _write_endpoint_state(api_key, enabled = True) + request.app.state.external_api_key = api_key + + local_url, external_url = _build_endpoint_urls(request) + return { + "enabled": True, + "api_key": api_key, + "local_url": local_url, + "external_url": external_url, + "model": _get_loaded_model_name(), + } + + @router.post("/chat/completions") async def openai_chat_completions( payload: ChatCompletionRequest, diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 4841c5d0a3..9148f3105f 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -50,6 +50,10 @@ def auth_db_path() -> Path: return auth_root() / "auth.db" +def access_endpoint_state_path() -> Path: + return auth_root() / "access_endpoint.json" + + def studio_db_path() -> Path: return studio_root() / "studio.db" diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 00d407fab9..a7a173c22f 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -692,6 +692,28 @@ print(completion.choices[0].message.content)`, } }, [endpointEnabled, applyEndpointData]); + const handleRegenerateKey = useCallback(async () => { + setEndpointLoading(true); + try { + const res = await authFetch( + "/api/inference/access-endpoint/regenerate", + { method: "POST" }, + ); + if (!res.ok) { + const err = await res.json().catch(() => null); + toast.error(err?.detail ?? "Failed to regenerate API key"); + return; + } + const data = await res.json(); + applyEndpointData(data); + toast.success("API key regenerated"); + } catch { + toast.error("Failed to regenerate API key"); + } finally { + setEndpointLoading(false); + } + }, [applyEndpointData]); + const handleCheckpointChange = useCallback( ( value: string, @@ -1247,12 +1269,24 @@ print(completion.choices[0].message.content)`, - +
+ + +