Use HF username for personal trace uploads (#199)

* Use HF username for personal trace uploads

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

* Remove redundant HF token branch

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

---------

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

agent/config.py

@@ -31,7 +31,7 @@ class Config(BaseModel):
     # format so the HF Agent Trace Viewer auto-renders it
     # (https://huggingface.co/changelog/agent-trace-viewer). Created private
     # on first use; user flips it public via /share-traces. ``{hf_user}`` is
-    # substituted at upload time from ``Session.user_id``.
+    # substituted at upload time from the authenticated HF username.
     share_traces: bool = True
     personal_trace_repo_template: str = "{hf_user}/ml-intern-sessions"
     auto_save_interval: int = 1  # Save every N user turns (0 = disabled)

agent/core/session.py

@@ -89,10 +89,12 @@ class Session:
         defer_turn_complete_notification: bool = False,
         session_id: str | None = None,
         user_id: str | None = None,
+        hf_username: str | None = None,
         persistence_store: Any | None = None,
     ):
         self.hf_token: Optional[str] = hf_token
         self.user_id: Optional[str] = user_id
+        self.hf_username: Optional[str] = hf_username
         self.persistence_store = persistence_store
         self.tool_router = tool_router
         self.stream = stream
@@ -363,6 +365,7 @@ class Session:
         return {
             "session_id": self.session_id,
             "user_id": self.user_id,
+            "hf_username": self.hf_username,
             "session_start_time": self.session_start_time,
             "session_end_time": datetime.now().isoformat(),
             "model_name": self.config.model_name,
@@ -458,7 +461,7 @@ class Session:
             return False
 
     def _personal_trace_repo_id(self) -> Optional[str]:
-        """Resolve the per-user trace repo id from config + user_id.
+        """Resolve the per-user trace repo id from config + HF username.
 
         Returns ``None`` when sharing is disabled, the user is anonymous,
         or the template is missing — caller skips the personal upload in
@@ -466,13 +469,14 @@ class Session:
         """
         if not getattr(self.config, "share_traces", False):
             return None
-        if not self.user_id:
+        hf_user = self.hf_username or self.user_id
+        if not hf_user:
             return None
         template = getattr(self.config, "personal_trace_repo_template", None)
         if not template:
             return None
         try:
-            return template.format(hf_user=self.user_id)
+            return template.format(hf_user=hf_user)
         except (KeyError, IndexError):
             logger.debug("personal_trace_repo_template format failed: %r", template)
             return None

backend/routes/agent.py

@@ -185,6 +185,7 @@ async def _check_session_access(
         session_id,
         user["user_id"],
         hf_token=hf_token,
+        hf_username=user.get("username"),
     )
     if not agent_session:
         raise HTTPException(status_code=404, detail="Session not found")
@@ -369,6 +370,7 @@ async def create_session(
     try:
         session_id = await session_manager.create_session(
             user_id=user["user_id"],
+            hf_username=user.get("username"),
             hf_token=hf_token,
             model=model,
             is_pro=user.get("plan") == "pro",
@@ -408,6 +410,7 @@ async def restore_session_summary(
     try:
         session_id = await session_manager.create_session(
             user_id=user["user_id"],
+            hf_username=user.get("username"),
             hf_token=hf_token,
             model=model,
             is_pro=user.get("plan") == "pro",

backend/session_manager.py

@@ -87,6 +87,7 @@ class AgentSession:
     tool_router: ToolRouter
     submission_queue: asyncio.Queue
     user_id: str = "dev"  # Owner of this session
+    hf_username: str | None = None  # HF namespace used for personal trace uploads
     hf_token: str | None = None  # User's HF OAuth token for tool execution
     task: asyncio.Task | None = None
     created_at: datetime = field(default_factory=datetime.utcnow)
@@ -157,6 +158,7 @@ class SessionManager:
         *,
         session_id: str,
         user_id: str,
+        hf_username: str | None,
         hf_token: str | None,
         model: str | None,
         event_queue: asyncio.Queue,
@@ -178,6 +180,7 @@ class SessionManager:
             tool_router=tool_router,
             hf_token=hf_token,
             user_id=user_id,
+            hf_username=hf_username,
             notification_gateway=self.messaging_gateway,
             notification_destinations=notification_destinations or [],
             session_id=session_id,
@@ -327,11 +330,18 @@ class SessionManager:
         )
 
     @staticmethod
-    def _update_hf_token(agent_session: AgentSession, hf_token: str | None) -> None:
-        if not hf_token:
-            return
-        agent_session.hf_token = hf_token
-        agent_session.session.hf_token = hf_token
+    def _update_hf_identity(
+        agent_session: AgentSession,
+        *,
+        hf_token: str | None,
+        hf_username: str | None,
+    ) -> None:
+        if hf_token:
+            agent_session.hf_token = hf_token
+            agent_session.session.hf_token = hf_token
+        if hf_username:
+            agent_session.hf_username = hf_username
+            agent_session.session.hf_username = hf_username
 
     async def persist_session_snapshot(
         self,
@@ -373,13 +383,18 @@ class SessionManager:
         session_id: str,
         user_id: str,
         hf_token: str | None = None,
+        hf_username: str | None = None,
     ) -> AgentSession | None:
         """Return a live runtime session, lazily restoring it from Mongo."""
         async with self._lock:
             existing = self.sessions.get(session_id)
         if existing:
             if self._can_access_session(existing, user_id):
-                self._update_hf_token(existing, hf_token)
+                self._update_hf_identity(
+                    existing,
+                    hf_token=hf_token,
+                    hf_username=hf_username,
+                )
                 return existing
             return None
 
@@ -392,7 +407,11 @@ class SessionManager:
             existing = self.sessions.get(session_id)
         if existing:
             if self._can_access_session(existing, user_id):
-                self._update_hf_token(existing, hf_token)
+                self._update_hf_identity(
+                    existing,
+                    hf_token=hf_token,
+                    hf_username=hf_username,
+                )
                 return existing
             return None
 
@@ -410,6 +429,7 @@ class SessionManager:
             self._create_session_sync,
             session_id=session_id,
             user_id=owner or user_id,
+            hf_username=hf_username,
             hf_token=hf_token,
             model=model,
             event_queue=event_queue,
@@ -442,6 +462,7 @@ class SessionManager:
             tool_router=tool_router,
             submission_queue=submission_queue,
             user_id=owner or user_id,
+            hf_username=hf_username,
             hf_token=hf_token,
             created_at=created_at,
             is_active=True,
@@ -455,7 +476,11 @@ class SessionManager:
             tool_router=tool_router,
         )
         if started is not agent_session:
-            self._update_hf_token(started, hf_token)
+            self._update_hf_identity(
+                started,
+                hf_token=hf_token,
+                hf_username=hf_username,
+            )
             return started
         logger.info("Restored session %s for user %s", session_id, owner or user_id)
         return agent_session
@@ -463,6 +488,7 @@ class SessionManager:
     async def create_session(
         self,
         user_id: str = "dev",
+        hf_username: str | None = None,
         hf_token: str | None = None,
         model: str | None = None,
         is_pro: bool | None = None,
@@ -475,6 +501,7 @@ class SessionManager:
 
         Args:
             user_id: The ID of the user who owns this session.
+            hf_username: The HF username/namespace used for personal trace uploads.
             hf_token: The user's HF OAuth token, stored for tool execution.
             model: Optional model override. When set, replaces ``model_name``
                 on the per-session config clone. None falls back to the
@@ -513,6 +540,7 @@ class SessionManager:
             self._create_session_sync,
             session_id=session_id,
             user_id=user_id,
+            hf_username=hf_username,
             hf_token=hf_token,
             model=model,
             event_queue=event_queue,
@@ -525,6 +553,7 @@ class SessionManager:
             tool_router=tool_router,
             submission_queue=submission_queue,
             user_id=user_id,
+            hf_username=hf_username,
             hf_token=hf_token,
         )
 

tests/unit/test_personal_trace_repo.py

@@ -0,0 +1,43 @@
+import asyncio
+from types import SimpleNamespace
+
+from agent.core.session import Session
+
+
+class DummyToolRouter:
+    def get_tool_specs_for_llm(self) -> list[dict]:
+        return []
+
+
+def _session(*, user_id: str | None, hf_username: str | None) -> Session:
+    config = SimpleNamespace(
+        model_name="moonshotai/Kimi-K2.6",
+        save_sessions=True,
+        share_traces=True,
+        personal_trace_repo_template="{hf_user}/ml-intern-sessions",
+        session_dataset_repo="smolagents/ml-intern-sessions",
+        auto_save_interval=1,
+        heartbeat_interval_s=0,
+        reasoning_effort=None,
+    )
+    context_manager = SimpleNamespace(items=[], on_message_added=None)
+    return Session(
+        event_queue=asyncio.Queue(),
+        config=config,
+        tool_router=DummyToolRouter(),
+        context_manager=context_manager,
+        user_id=user_id,
+        hf_username=hf_username,
+    )
+
+
+def test_personal_trace_repo_uses_hf_username_before_oauth_subject():
+    session = _session(user_id="oauth-subject", hf_username="lewtun")
+
+    assert session._personal_trace_repo_id() == "lewtun/ml-intern-sessions"
+
+
+def test_personal_trace_repo_falls_back_to_user_id_for_cli():
+    session = _session(user_id="lewtun", hf_username=None)
+
+    assert session._personal_trace_repo_id() == "lewtun/ml-intern-sessions"