Use `huggingface_hub.get_token()` for auth fallback (#126)

* Use cached HF login token for router auth fallback

Let HF router requests resolve credentials from huggingface_hub's local login cache when INFERENCE_TOKEN, session token, and HF_TOKEN are absent. This keeps CLI and plugin flows working after `hf auth login` without requiring explicit env exports.

Made-with: Cursor

* Consolidate Hugging Face token resolution

Centralize HF token cleanup and lookup so router calls, CLI startup, and backend request paths do not drift after adding huggingface_hub.get_token() fallback. The backend request helper intentionally keeps browser/user token precedence separate from local cached CLI auth.

Constraint: huggingface_hub.get_token() already handles HF_TOKEN, HUGGING_FACE_HUB_TOKEN, and the local hf auth login cache.

Rejected: Use cached login tokens for backend request auth | would let server-local credentials stand in for the browser user.

Confidence: high

Scope-risk: narrow

Directive: Keep router-token fallback and request-user token extraction separate; only router/CLI paths should use cached login fallback.

Tested: UV_CACHE_DIR=/tmp/uv-cache uv run --extra dev pytest tests/unit/test_llm_params.py -q

Tested: UV_CACHE_DIR=/tmp/uv-cache uv run --extra dev pytest tests/unit/test_hf_access.py tests/unit/test_user_quotas.py -q

Tested: python -m py_compile agent/core/hf_tokens.py agent/core/llm_params.py agent/main.py backend/routes/agent.py backend/dependencies.py

Tested: unset HF_TOKEN/INFERENCE_TOKEN cached-token live HfApi.whoami smoke

* Preserve session user id while deduping HF token lookup

The token helper cleanup should not change session identity plumbing. Restore the existing username helper and pass user_id into both interactive and headless submission_loop calls, while removing only the redundant _get_hf_token wrapper.

Constraint: PR #146 relies on user_id being passed into Session for saved-session ownership.

Rejected: Inline HfApi.whoami at each call site | reintroduces duplicate username lookup code and obscures the token-only cleanup.

Confidence: high

Scope-risk: narrow

Directive: Do not remove user_id from submission_loop call sites when editing CLI auth/token code.

Tested: python -m py_compile agent/main.py agent/core/hf_tokens.py agent/core/llm_params.py backend/routes/agent.py backend/dependencies.py

Tested: UV_CACHE_DIR=/tmp/uv-cache uv run --extra dev pytest tests/unit/test_llm_params.py -q

Tested: UV_CACHE_DIR=/tmp/uv-cache uv run --extra dev pytest tests/unit/test_cli_rendering.py tests/unit/test_hf_access.py tests/unit/test_user_quotas.py -q

---------

Co-authored-by: Aksel Joonas Reedi <125026660+akseljoonas@users.noreply.github.com>
Co-authored-by: akseljoonas <aksel.reedi@gmail.com>

agent/core/hf_tokens.py

@@ -0,0 +1,85 @@
+"""Hugging Face token resolution helpers."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+
+def clean_hf_token(token: str | None) -> str | None:
+    """Normalize token strings the same way huggingface_hub does."""
+    if token is None:
+        return None
+    return token.replace("\r", "").replace("\n", "").strip() or None
+
+
+def get_cached_hf_token() -> str | None:
+    """Return the token from huggingface_hub's normal env/cache lookup."""
+    try:
+        from huggingface_hub import get_token
+
+        return get_token()
+    except Exception:
+        return None
+
+
+def resolve_hf_token(
+    *candidates: str | None,
+    include_cached: bool = True,
+) -> str | None:
+    """Return the first non-empty explicit token, then optionally HF cache."""
+    for token in candidates:
+        cleaned = clean_hf_token(token)
+        if cleaned:
+            return cleaned
+    if include_cached:
+        return get_cached_hf_token()
+    return None
+
+
+def resolve_hf_router_token(session_hf_token: str | None = None) -> str | None:
+    """Resolve the token used for Hugging Face Router LLM calls.
+
+    App-specific precedence:
+    1. INFERENCE_TOKEN: shared hosted-Space inference token.
+    2. session_hf_token: the active user/session token.
+    3. huggingface_hub.get_token(): HF_TOKEN/HUGGING_FACE_HUB_TOKEN or
+       local ``hf auth login`` cache.
+    """
+    return resolve_hf_token(os.environ.get("INFERENCE_TOKEN"), session_hf_token)
+
+
+def get_hf_bill_to() -> str | None:
+    """Return X-HF-Bill-To only when a shared inference token is active."""
+    if clean_hf_token(os.environ.get("INFERENCE_TOKEN")):
+        return os.environ.get("HF_BILL_TO", "smolagents")
+    return None
+
+
+def bearer_token_from_header(auth_header: str | None) -> str | None:
+    """Extract a cleaned bearer token from an Authorization header."""
+    if not auth_header or not auth_header.startswith("Bearer "):
+        return None
+    return clean_hf_token(auth_header[7:])
+
+
+def resolve_hf_request_token(
+    request: Any,
+    *,
+    include_env_fallback: bool = True,
+) -> str | None:
+    """Resolve a user token from a FastAPI request.
+
+    This intentionally does not use the local ``hf auth login`` cache. Backend
+    request paths should act as the browser user from Authorization/cookie, or
+    fall back only to an explicit server ``HF_TOKEN`` in dev/server contexts.
+    """
+    token = bearer_token_from_header(request.headers.get("Authorization", ""))
+    if token:
+        return token
+    token = clean_hf_token(request.cookies.get("hf_access_token"))
+    if token:
+        return token
+    if include_env_fallback:
+        return clean_hf_token(os.environ.get("HF_TOKEN"))
+    return None

agent/core/llm_params.py

@@ -5,7 +5,12 @@ can import it without pulling in the whole agent loop / tool router and
 creating circular imports.
 """
 
-import os
+from agent.core.hf_tokens import get_hf_bill_to, resolve_hf_router_token
+
+
+def _resolve_hf_router_token(session_hf_token: str | None = None) -> str | None:
+    """Backward-compatible private wrapper used by tests and older imports."""
+    return resolve_hf_router_token(session_hf_token)
 
 
 def _patch_litellm_effort_validation() -> None:
@@ -129,7 +134,8 @@ def _resolve_llm_params(
       1. INFERENCE_TOKEN env — shared key on the hosted Space (inference is
          free for users, billed to the Space owner via ``X-HF-Bill-To``).
       2. session.hf_token — the user's own token (CLI / OAuth / cache file).
-      3. HF_TOKEN env — belt-and-suspenders fallback for CLI users.
+      3. huggingface_hub cache — ``HF_TOKEN`` / ``HUGGING_FACE_HUB_TOKEN`` /
+         local ``hf auth login`` cache.
     """
     if model_name.startswith("anthropic/"):
         params: dict = {"model": model_name}
@@ -175,18 +181,13 @@ def _resolve_llm_params(
         return params
 
     hf_model = model_name.removeprefix("huggingface/")
-    api_key = (
-        os.environ.get("INFERENCE_TOKEN")
-        or session_hf_token
-        or os.environ.get("HF_TOKEN")
-    )
+    api_key = _resolve_hf_router_token(session_hf_token)
     params = {
         "model": f"openai/{hf_model}",
         "api_base": "https://router.huggingface.co/v1",
         "api_key": api_key,
     }
-    if os.environ.get("INFERENCE_TOKEN"):
-        bill_to = os.environ.get("HF_BILL_TO", "smolagents")
+    if bill_to := get_hf_bill_to():
         params["extra_headers"] = {"X-HF-Bill-To": bill_to}
     if reasoning_effort:
         hf_level = "low" if reasoning_effort == "minimal" else reasoning_effort

agent/main.py

@@ -23,6 +23,7 @@ from prompt_toolkit import PromptSession
 from agent.config import load_config
 from agent.core.agent_loop import submission_loop
 from agent.core import model_switcher
+from agent.core.hf_tokens import resolve_hf_token
 from agent.core.session import OpType
 from agent.core.tools import ToolRouter
 from agent.utils.reliability_checks import check_training_script_save_pattern
@@ -69,28 +70,6 @@ def _safe_get_args(arguments: dict) -> dict:
     return args if isinstance(args, dict) else {}
 
 
-def _get_hf_token() -> str | None:
-    """Get HF token from environment, huggingface_hub API, or cached token file."""
-    token = os.environ.get("HF_TOKEN")
-    if token:
-        return token
-    try:
-        from huggingface_hub import HfApi
-        api = HfApi()
-        token = api.token
-        if token:
-            return token
-    except Exception:
-        pass
-    # Fallback: read the cached token file directly
-    token_path = Path.home() / ".cache" / "huggingface" / "token"
-    if token_path.exists():
-        token = token_path.read_text().strip()
-        if token:
-            return token
-    return None
-
-
 def _get_hf_user(token: str | None) -> str | None:
     """Resolve the HF username for a token, if available."""
     if not token:
@@ -769,7 +748,7 @@ async def _handle_slash_command(
         normalized = arg.removeprefix("huggingface/")
         session = session_holder[0] if session_holder else None
         await model_switcher.probe_and_switch_model(
-            normalized, config, session, console, _get_hf_token(),
+            normalized, config, session, console, resolve_hf_token(),
         )
         return None
 
@@ -838,7 +817,7 @@ async def main():
     prompt_session = PromptSession()
 
     # HF token — required, prompt if missing
-    hf_token = _get_hf_token()
+    hf_token = resolve_hf_token()
     if not hf_token:
         hf_token = await _prompt_and_save_hf_token(prompt_session)
 
@@ -1054,7 +1033,7 @@ async def headless_main(
     logging.basicConfig(level=logging.WARNING)
     _configure_runtime_logging()
 
-    hf_token = _get_hf_token()
+    hf_token = resolve_hf_token()
     if not hf_token:
         print("ERROR: No HF token found. Set HF_TOKEN or run `huggingface-cli login`.", file=sys.stderr)
         sys.exit(1)

backend/dependencies.py

@@ -12,6 +12,8 @@ from typing import Any
 import httpx
 from fastapi import HTTPException, Request, status
 
+from agent.core.hf_tokens import bearer_token_from_header
+
 from agent.core.hf_access import fetch_whoami_v2, jobs_access_from_whoami
 
 logger = logging.getLogger(__name__)
@@ -157,9 +159,8 @@ async def get_current_user(request: Request) -> dict[str, Any]:
         return DEV_USER
 
     # Try Authorization header
-    auth_header = request.headers.get("Authorization", "")
-    if auth_header.startswith("Bearer "):
-        token = auth_header[7:]
+    token = bearer_token_from_header(request.headers.get("Authorization", ""))
+    if token:
         user = await _extract_user_from_token(token)
         if user:
             return user
@@ -183,9 +184,9 @@ def _extract_token(request: Request) -> str | None:
 
     Mirrors the lookup order used by ``get_current_user``.
     """
-    auth_header = request.headers.get("Authorization", "")
-    if auth_header.startswith("Bearer "):
-        return auth_header[7:]
+    token = bearer_token_from_header(request.headers.get("Authorization", ""))
+    if token:
+        return token
     return request.cookies.get("hf_access_token")
 
 
@@ -202,4 +203,3 @@ async def require_huggingface_org_member(request: Request) -> bool:
     if not token:
         return False
     return await check_org_membership(token, HF_EMPLOYEE_ORG)
-

backend/routes/agent.py

@@ -33,6 +33,7 @@ from session_manager import MAX_SESSIONS, AgentSession, SessionCapacityError, se
 import user_quotas
 
 from agent.core.hf_access import get_jobs_access
+from agent.core.hf_tokens import resolve_hf_request_token, resolve_hf_router_token
 from agent.core.llm_params import _resolve_llm_params
 
 logger = logging.getLogger(__name__)
@@ -332,10 +333,8 @@ async def generate_title(
     reasoning model — reasoning_effort=low keeps the reasoning budget small
     so the 60-token output budget isn't consumed before the title is written.
     """
-    api_key = (
-        os.environ.get("INFERENCE_TOKEN")
-        or (user.get("hf_token") if isinstance(user, dict) else None)
-        or os.environ.get("HF_TOKEN")
+    api_key = resolve_hf_router_token(
+        user.get("hf_token") if isinstance(user, dict) else None
     )
     try:
         response = await acompletion(
@@ -391,14 +390,7 @@ async def create_session(
     Returns 503 if the server or user has reached the session limit.
     """
     # Extract the user's HF token (Bearer header, HttpOnly cookie, or env var)
-    hf_token = None
-    auth_header = request.headers.get("Authorization", "")
-    if auth_header.startswith("Bearer "):
-        hf_token = auth_header[7:]
-    if not hf_token:
-        hf_token = request.cookies.get("hf_access_token")
-    if not hf_token:
-        hf_token = os.environ.get("HF_TOKEN")
+    hf_token = resolve_hf_request_token(request)
 
     # Optional model override. Empty body falls back to the config default.
     model: str | None = None
@@ -444,14 +436,7 @@ async def restore_session_summary(
     if not isinstance(messages, list) or not messages:
         raise HTTPException(status_code=400, detail="Missing 'messages' array")
 
-    hf_token = None
-    auth_header = request.headers.get("Authorization", "")
-    if auth_header.startswith("Bearer "):
-        hf_token = auth_header[7:]
-    if not hf_token:
-        hf_token = request.cookies.get("hf_access_token")
-    if not hf_token:
-        hf_token = os.environ.get("HF_TOKEN")
+    hf_token = resolve_hf_request_token(request)
 
     model = body.get("model")
     valid_ids = {m["id"] for m in AVAILABLE_MODELS}
@@ -545,14 +530,7 @@ async def get_user_quota(user: dict = Depends(get_current_user)) -> dict:
 @router.get("/user/jobs-access")
 async def get_jobs_access_info(request: Request, user: dict = Depends(get_current_user)) -> dict:
     """Return whether the current token can run HF Jobs and under which namespaces."""
-    token = None
-    auth_header = request.headers.get("Authorization", "")
-    if auth_header.startswith("Bearer "):
-        token = auth_header[7:]
-    if not token:
-        token = request.cookies.get("hf_access_token")
-    if not token:
-        token = os.environ.get("HF_TOKEN")
+    token = resolve_hf_request_token(request)
 
     access = await get_jobs_access(token or "")
     return {

tests/unit/test_llm_params.py

@@ -1,4 +1,9 @@
-from agent.core.llm_params import UnsupportedEffortError, _resolve_llm_params
+from agent.core.hf_tokens import resolve_hf_request_token
+from agent.core.llm_params import (
+    UnsupportedEffortError,
+    _resolve_hf_router_token,
+    _resolve_llm_params,
+)
 
 
 def test_openai_xhigh_effort_is_forwarded():
@@ -23,3 +28,80 @@ def test_openai_max_effort_is_still_rejected():
         assert "OpenAI doesn't accept effort='max'" in str(exc)
     else:
         raise AssertionError("Expected UnsupportedEffortError for max effort")
+
+
+def test_hf_router_token_prefers_inference_token(monkeypatch):
+    monkeypatch.setenv("INFERENCE_TOKEN", " inference-token ")
+    monkeypatch.setenv("HF_TOKEN", "hf-token")
+
+    assert _resolve_hf_router_token("session-token") == "inference-token"
+
+
+def test_hf_router_token_prefers_session_over_hf_cache(monkeypatch):
+    monkeypatch.delenv("INFERENCE_TOKEN", raising=False)
+    monkeypatch.setenv("HF_TOKEN", "hf-token")
+
+    assert _resolve_hf_router_token(" session-token ") == "session-token"
+
+
+def test_hf_router_token_uses_hf_token_env_via_huggingface_hub(monkeypatch):
+    monkeypatch.delenv("INFERENCE_TOKEN", raising=False)
+    monkeypatch.setenv("HF_TOKEN", " hf-token ")
+
+    assert _resolve_hf_router_token(None) == "hf-token"
+
+
+def test_hf_router_token_uses_huggingface_hub_cache(monkeypatch):
+    import huggingface_hub
+
+    monkeypatch.delenv("INFERENCE_TOKEN", raising=False)
+    monkeypatch.delenv("HF_TOKEN", raising=False)
+    monkeypatch.setattr(huggingface_hub, "get_token", lambda: "cached-token")
+
+    assert _resolve_hf_router_token(None) == "cached-token"
+
+
+def test_hf_router_token_swallows_huggingface_hub_errors(monkeypatch):
+    import huggingface_hub
+
+    def fail():
+        raise RuntimeError("cache unavailable")
+
+    monkeypatch.delenv("INFERENCE_TOKEN", raising=False)
+    monkeypatch.delenv("HF_TOKEN", raising=False)
+    monkeypatch.setattr(huggingface_hub, "get_token", fail)
+
+    assert _resolve_hf_router_token(None) is None
+
+
+def test_hf_router_params_set_bill_to_only_for_inference_token(monkeypatch):
+    monkeypatch.setenv("INFERENCE_TOKEN", "inference-token")
+    monkeypatch.setenv("HF_BILL_TO", "test-org")
+
+    params = _resolve_llm_params("moonshotai/Kimi-K2.6")
+
+    assert params["api_key"] == "inference-token"
+    assert params["extra_headers"] == {"X-HF-Bill-To": "test-org"}
+
+
+def test_hf_request_token_keeps_browser_user_precedence(monkeypatch):
+    class Request:
+        headers = {"Authorization": "Bearer browser-token"}
+        cookies = {"hf_access_token": "cookie-token"}
+
+    monkeypatch.setenv("HF_TOKEN", "server-token")
+
+    assert resolve_hf_request_token(Request()) == "browser-token"
+
+
+def test_hf_request_token_does_not_use_cached_login(monkeypatch):
+    import huggingface_hub
+
+    class Request:
+        headers = {}
+        cookies = {}
+
+    monkeypatch.delenv("HF_TOKEN", raising=False)
+    monkeypatch.setattr(huggingface_hub, "get_token", lambda: "cached-token")
+
+    assert resolve_hf_request_token(Request()) is None