Fall back to user HF token for router when INFERENCE_TOKEN is unset (#39)

akseljoonas · claude · web-flow · commit 182ddeef0cdc · 2026-04-21T20:19:16.000+03:00
The HF router code path only read INFERENCE_TOKEN, which is the shared server-side key set on the hosted Space so inference is free for users. On the CLI / self-hosted path that env var is absent, so requests went out with no bearer token and the router returned 401 — surfaced to users as "Authentication failed" even with a valid HF_TOKEN (issue #36). Resolve api_key in this order: 1. INFERENCE_TOKEN env (unchanged Space behavior — shared billing) 2. session.hf_token (user's OAuth / CLI token) 3. HF_TOKEN env (belt-and-suspenders for CLI) Applied to _resolve_hf_router_params, research_tool._resolve_llm_params, and ContextManager.compact. Fixes #36 Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
diff --git a/agent/context_manager/manager.py b/agent/context_manager/manager.py
@@ -263,7 +263,10 @@ def truncate_to_user_message(self, user_message_index: int) -> bool:
         return False
 
     async def compact(
-        self, model_name: str, tool_specs: list[dict] | None = None
+        self,
+        model_name: str,
+        tool_specs: list[dict] | None = None,
+        hf_token: str | None = None,
     ) -> None:
         """Remove old messages to keep history under target size"""
         if (self.context_length <= self.max_context) or not self.items:
@@ -303,7 +306,11 @@ async def compact(
             )
         )
 
-        hf_key = os.environ.get("INFERENCE_TOKEN")
+        hf_key = (
+            os.environ.get("INFERENCE_TOKEN")
+            or hf_token
+            or os.environ.get("HF_TOKEN")
+        )
         response = await acompletion(
             model=model_name,
             messages=messages_to_summarize,
diff --git a/agent/core/agent_loop.py b/agent/core/agent_loop.py
@@ -20,11 +20,11 @@
 logger = logging.getLogger(__name__)
 
 ToolCall = ChatCompletionMessageToolCall
-# Explicit inference token for LLM API calls (separate from user OAuth tokens).
-_INFERENCE_API_KEY = os.environ.get("INFERENCE_TOKEN")
 
 
-def _resolve_hf_router_params(model_name: str) -> dict:
+def _resolve_hf_router_params(
+    model_name: str, session_hf_token: str | None = None
+) -> dict:
     """
     Build LiteLLM kwargs for HuggingFace Router models.
 
@@ -35,6 +35,13 @@ def _resolve_hf_router_params(model_name: str) -> dict:
 
     Input format:  huggingface/<router_provider>/<org>/<model>
     Example:       huggingface/novita/moonshotai/kimi-k2.5
+
+    Token resolution (first non-empty wins):
+      1. INFERENCE_TOKEN env — shared key on the hosted Space so inference
+         is free for users and billed to the Space owner.
+      2. session.hf_token — the user's own token (CLI or self-hosted),
+         resolved from env / huggingface-cli login / cached token file.
+      3. HF_TOKEN env — belt-and-suspenders fallback for CLI users.
     """
     if not model_name.startswith("huggingface/"):
         return {"model": model_name}
@@ -47,7 +54,11 @@ def _resolve_hf_router_params(model_name: str) -> dict:
 
     router_provider = parts[1]
     actual_model = parts[2]
-    api_key = _INFERENCE_API_KEY
+    api_key = (
+        os.environ.get("INFERENCE_TOKEN")
+        or session_hf_token
+        or os.environ.get("HF_TOKEN")
+    )
 
     return {
         "model": f"openai/{actual_model}",
@@ -205,6 +216,7 @@ async def _compact_and_notify(session: Session) -> None:
     await session.context_manager.compact(
         model_name=session.config.model_name,
         tool_specs=tool_specs,
+        hf_token=session.hf_token,
     )
     new_length = session.context_manager.context_length
     if new_length != old_length:
@@ -506,7 +518,9 @@ async def run_agent(
             tools = session.tool_router.get_tool_specs_for_llm()
             try:
                 # ── Call the LLM (streaming or non-streaming) ──
-                llm_params = _resolve_hf_router_params(session.config.model_name)
+                llm_params = _resolve_hf_router_params(
+                    session.config.model_name, session.hf_token
+                )
                 if session.stream:
                     llm_result = await _call_llm_streaming(session, messages, tools, llm_params)
                 else:
diff --git a/agent/tools/research_tool.py b/agent/tools/research_tool.py
@@ -213,7 +213,9 @@
 }
 
 
-def _resolve_llm_params(model_name: str) -> dict:
+def _resolve_llm_params(
+    model_name: str, session_hf_token: str | None = None
+) -> dict:
     """Build LiteLLM kwargs, reusing the HF router logic from agent_loop."""
     if not model_name.startswith("huggingface/"):
         return {"model": model_name}
@@ -224,10 +226,16 @@ def _resolve_llm_params(model_name: str) -> dict:
 
     provider = parts[1]
     model_id = parts[2]
+    api_key = (
+        os.environ.get("INFERENCE_TOKEN")
+        or session_hf_token
+        or os.environ.get("HF_TOKEN")
+        or ""
+    )
     return {
         "model": f"openai/{model_id}",
         "api_base": f"https://router.huggingface.co/{provider}/v3/openai",
-        "api_key": os.environ.get("INFERENCE_TOKEN", ""),
+        "api_key": api_key,
     }
 
 
@@ -264,7 +272,7 @@ async def research_handler(
     # Use a cheaper/faster model for research
     main_model = session.config.model_name
     research_model = _get_research_model(main_model)
-    llm_params = _resolve_llm_params(research_model)
+    llm_params = _resolve_llm_params(research_model, getattr(session, "hf_token", None))
 
     # Get read-only tool specs from the session's tool router
     tool_specs = [
