ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: persistent Access Endpoint API key with Regenerate button

Browse source 
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.
This commit is contained in:
Roland Tannous 2026-04-09 20:24:19 +00:00
commit 8a574de915
4 changed files with 151 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

21
studio/backend/main.py
View file

@ -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

100
studio/backend/routes/inference.py
View file

@ -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,

4
studio/backend/utils/paths/storage_roots.py
View file

@ -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"

46
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -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)`,
<label htmlFor="endpoint-api-key" className="text-xs font-medium">
API Key
</label>
<Input
id="endpoint-api-key"
value={endpointApiKey}
readOnly
className="font-mono text-xs"
/>
<div className="flex items-center gap-2">
<Input
id="endpoint-api-key"
value={endpointApiKey}
readOnly
className="font-mono text-xs"
/>
<Button
variant="ghost"
size="sm"
className="h-8 shrink-0 text-xs"
onClick={handleRegenerateKey}
disabled={endpointLoading}
title="Generate a new API key. The previous key will stop working."
>
Regenerate
</Button>
</div>
</div>
<div className="grid gap-1.5">
<label htmlFor="endpoint-model-alias" className="text-xs font-medium">