Deploy 2026-05-04

Co-authored-by: OpenAI Codex <codex@openai.com>

backend/dependencies.py

@@ -12,7 +12,7 @@ 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_tokens import bearer_token_from_header, clean_hf_token
 
 from agent.core.hf_access import fetch_whoami_v2
 
@@ -36,6 +36,8 @@ DEV_USER: dict[str, Any] = {
     "plan": "org",  # Dev runs at the Pro/Org quota tier so local testing isn't capped.
 }
 
+INTERNAL_HF_TOKEN_KEY = "_hf_token"
+
 # Plan field discovery — log the whoami-v2 shape once at DEBUG so we can
 # confirm the actual key in production without hammering the HF API.
 _WHOAMI_SHAPE_LOGGED = False
@@ -135,6 +137,7 @@ async def _extract_user_from_token(token: str) -> dict[str, Any] | None:
         return None
     user = _user_from_info(user_info)
     user["plan"] = await _fetch_user_plan(token)
+    user[INTERNAL_HF_TOKEN_KEY] = clean_hf_token(token)
     return user
 
 
@@ -145,13 +148,13 @@ async def _dev_user_from_env() -> dict[str, Any]:
     real HF namespace. Deriving the dev user from HF_TOKEN keeps local uploads
     pointed at the token owner's dataset instead of dev/ml-intern-sessions.
     """
-    token = os.environ.get("HF_TOKEN", "")
+    token = clean_hf_token(os.environ.get("HF_TOKEN", ""))
     if not token:
-        return DEV_USER
+        return dict(DEV_USER)
 
     whoami = await fetch_whoami_v2(token)
     if not isinstance(whoami, dict):
-        return DEV_USER
+        return dict(DEV_USER)
 
     username = None
     for key in ("name", "user", "preferred_username"):
@@ -160,13 +163,14 @@ async def _dev_user_from_env() -> dict[str, Any]:
             username = value
             break
     if not username:
-        return DEV_USER
+        return dict(DEV_USER)
 
     return {
         "user_id": username,
         "username": username,
         "authenticated": True,
         "plan": await _fetch_user_plan(token),
+        INTERNAL_HF_TOKEN_KEY: token,
     }
 
 

backend/routes/agent.py

@@ -10,7 +10,11 @@ import logging
 import os
 from typing import Any
 
-from dependencies import get_current_user, require_huggingface_org_member
+from dependencies import (
+    INTERNAL_HF_TOKEN_KEY,
+    get_current_user,
+    require_huggingface_org_member,
+)
 from fastapi import (
     APIRouter,
     Depends,
@@ -209,6 +213,12 @@ async def _enforce_gated_model_quota(
     await session_manager.persist_session_snapshot(agent_session)
 
 
+def _user_hf_token(user: dict[str, Any] | None) -> str | None:
+    if not isinstance(user, dict):
+        return None
+    return user.get(INTERNAL_HF_TOKEN_KEY)
+
+
 async def _check_session_access(
     session_id: str,
     user: dict[str, Any],
@@ -216,7 +226,7 @@ async def _check_session_access(
     preload_sandbox: bool = True,
 ) -> AgentSession:
     """Verify and lazily load the user's session. Raises 403 or 404."""
-    hf_token = resolve_hf_request_token(request) if request is not None else user.get("hf_token")
+    hf_token = resolve_hf_request_token(request) if request is not None else _user_hf_token(user)
     agent_session = await session_manager.ensure_session_loaded(
         session_id,
         user["user_id"],
@@ -318,9 +328,7 @@ 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 = resolve_hf_router_token(
-        user.get("hf_token") if isinstance(user, dict) else None
-    )
+    api_key = resolve_hf_router_token(_user_hf_token(user))
     try:
         response = await acompletion(
             # Double openai/ prefix: LiteLLM strips the first as its provider

backend/routes/auth.py

@@ -167,6 +167,5 @@ async def get_me(user: dict = Depends(get_current_user)) -> dict:
 
     Uses the shared auth dependency which handles cookie + Bearer token.
     """
-    return user
-
+    return {key: value for key, value in user.items() if not key.startswith("_")}
 

backend/session_manager.py

@@ -376,6 +376,40 @@ class SessionManager:
             agent_session.hf_username = hf_username
             agent_session.session.hf_username = hf_username
 
+    @staticmethod
+    def _has_active_sandbox_preload(agent_session: AgentSession) -> bool:
+        task = getattr(agent_session.session, "sandbox_preload_task", None)
+        return bool(task and not task.done())
+
+    @staticmethod
+    def _preload_failed_for_missing_hf_token(agent_session: AgentSession) -> bool:
+        error = getattr(agent_session.session, "sandbox_preload_error", None)
+        return isinstance(error, str) and error.startswith("No HF token available")
+
+    def _restart_cpu_preload_if_token_recovered(
+        self,
+        agent_session: AgentSession,
+        *,
+        preload_sandbox: bool,
+    ) -> None:
+        if not preload_sandbox:
+            return
+        session = agent_session.session
+        if getattr(session, "sandbox", None):
+            return
+        if self._has_active_sandbox_preload(agent_session):
+            return
+        if not (agent_session.hf_token or getattr(session, "hf_token", None)):
+            return
+
+        if not self._preload_failed_for_missing_hf_token(agent_session):
+            return
+
+        session.sandbox_preload_error = None
+        session.sandbox_preload_task = None
+        session.sandbox_preload_cancel_event = None
+        self._start_cpu_sandbox_preload(agent_session)
+
     async def _clear_persisted_sandbox_metadata(self, session_id: str) -> None:
         try:
             await self._store().update_session_fields(
@@ -526,6 +560,10 @@ class SessionManager:
                     hf_token=hf_token,
                     hf_username=hf_username,
                 )
+                self._restart_cpu_preload_if_token_recovered(
+                    existing,
+                    preload_sandbox=preload_sandbox,
+                )
                 return existing
             return None
 
@@ -543,6 +581,10 @@ class SessionManager:
                     hf_token=hf_token,
                     hf_username=hf_username,
                 )
+                self._restart_cpu_preload_if_token_recovered(
+                    existing,
+                    preload_sandbox=preload_sandbox,
+                )
                 return existing
             return None
 

tests/unit/test_auth_token_propagation.py

@@ -0,0 +1,61 @@
+"""Tests for authenticated HF token propagation through backend dependencies."""
+
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+_BACKEND_DIR = Path(__file__).resolve().parent.parent.parent / "backend"
+if str(_BACKEND_DIR) not in sys.path:
+    sys.path.insert(0, str(_BACKEND_DIR))
+
+import dependencies  # noqa: E402
+from routes import auth  # noqa: E402
+
+
+@pytest.mark.asyncio
+async def test_current_user_carries_internal_hf_token(monkeypatch):
+    monkeypatch.setattr(dependencies, "AUTH_ENABLED", True)
+    dependencies._token_cache.clear()
+
+    async def fake_validate_token(token):
+        assert token == "hf-user-token"
+        return {"sub": "user-id", "preferred_username": "alice"}
+
+    async def fake_fetch_user_plan(token):
+        assert token == "hf-user-token"
+        return "pro"
+
+    monkeypatch.setattr(dependencies, "_validate_token", fake_validate_token)
+    monkeypatch.setattr(dependencies, "_fetch_user_plan", fake_fetch_user_plan)
+
+    request = SimpleNamespace(
+        headers={"Authorization": "Bearer hf-user-token"},
+        cookies={},
+    )
+
+    user = await dependencies.get_current_user(request)
+
+    assert user["user_id"] == "user-id"
+    assert user["username"] == "alice"
+    assert user["plan"] == "pro"
+    assert user[dependencies.INTERNAL_HF_TOKEN_KEY] == "hf-user-token"
+
+
+@pytest.mark.asyncio
+async def test_auth_me_does_not_expose_internal_hf_token():
+    user = {
+        "user_id": "user-id",
+        "username": "alice",
+        "authenticated": True,
+        dependencies.INTERNAL_HF_TOKEN_KEY: "hf-user-token",
+    }
+
+    response = await auth.get_me(user)
+
+    assert response == {
+        "user_id": "user-id",
+        "username": "alice",
+        "authenticated": True,
+    }

tests/unit/test_session_manager_persistence.py

@@ -281,6 +281,101 @@ async def test_existing_session_updates_token_after_access_check():
     assert existing.session.hf_token == "new-token"
 
 
+@pytest.mark.asyncio
+async def test_existing_session_retries_preload_after_token_recovered():
+    manager = _manager_with_store(NoopSessionStore())
+    existing = _runtime_agent_session("s1", user_id="owner", hf_token=None)
+    done_task = asyncio.get_running_loop().create_future()
+    done_task.set_result(None)
+    existing.session.sandbox_preload_task = done_task
+    existing.session.sandbox_preload_error = (
+        "No HF token available. Cannot create sandbox."
+    )
+    manager.sessions["s1"] = existing
+    started: list[str] = []
+
+    def fake_start_cpu_sandbox_preload(agent_session):
+        started.append(agent_session.session_id)
+
+    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]
+
+    result = await manager.ensure_session_loaded(
+        "s1",
+        user_id="owner",
+        hf_token="new-token",
+    )
+
+    assert result is existing
+    assert existing.hf_token == "new-token"
+    assert existing.session.hf_token == "new-token"
+    assert existing.session.sandbox_preload_error is None
+    assert existing.session.sandbox_preload_task is None
+    assert started == ["s1"]
+
+
+@pytest.mark.asyncio
+async def test_existing_session_does_not_retry_preload_when_disabled():
+    manager = _manager_with_store(NoopSessionStore())
+    existing = _runtime_agent_session("s1", user_id="owner", hf_token=None)
+    done_task = asyncio.get_running_loop().create_future()
+    done_task.set_result(None)
+    existing.session.sandbox_preload_task = done_task
+    existing.session.sandbox_preload_error = (
+        "No HF token available. Cannot create sandbox."
+    )
+    manager.sessions["s1"] = existing
+    started: list[str] = []
+
+    def fake_start_cpu_sandbox_preload(agent_session):
+        started.append(agent_session.session_id)
+
+    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]
+
+    result = await manager.ensure_session_loaded(
+        "s1",
+        user_id="owner",
+        hf_token="new-token",
+        preload_sandbox=False,
+    )
+
+    assert result is existing
+    assert existing.hf_token == "new-token"
+    assert existing.session.hf_token == "new-token"
+    assert existing.session.sandbox_preload_error == (
+        "No HF token available. Cannot create sandbox."
+    )
+    assert started == []
+
+
+@pytest.mark.asyncio
+async def test_existing_session_does_not_restart_preload_after_teardown():
+    manager = _manager_with_store(NoopSessionStore())
+    existing = _runtime_agent_session("s1", user_id="owner", hf_token="token")
+    done_task = asyncio.get_running_loop().create_future()
+    done_task.set_result(None)
+    existing.session.sandbox = None
+    existing.session.sandbox_preload_task = done_task
+    existing.session.sandbox_preload_error = None
+    manager.sessions["s1"] = existing
+    started: list[str] = []
+
+    def fake_start_cpu_sandbox_preload(agent_session):
+        started.append(agent_session.session_id)
+
+    manager._start_cpu_sandbox_preload = fake_start_cpu_sandbox_preload  # type: ignore[method-assign]
+
+    result = await manager.ensure_session_loaded(
+        "s1",
+        user_id="owner",
+        hf_token="token",
+    )
+
+    assert result is existing
+    assert existing.session.sandbox_preload_task is done_task
+    assert existing.session.sandbox_preload_error is None
+    assert started == []
+
+
 @pytest.mark.asyncio
 async def test_concurrent_lazy_restore_starts_only_one_agent_task():
     store = RestoreStore(delay=0.01)