sync: org onboarding, model updates, sandbox cleanup, prompt fixes

agent/core/session.py

@@ -19,12 +19,14 @@ logger = logging.getLogger(__name__)
 # on network calls for certain providers (known litellm issue).
 _MAX_TOKENS_MAP: dict[str, int] = {
     # Anthropic
+    "anthropic/claude-opus-4-6": 200_000,
     "anthropic/claude-opus-4-5-20251101": 200_000,
     "anthropic/claude-sonnet-4-5-20250929": 200_000,
     "anthropic/claude-sonnet-4-20250514": 200_000,
     "anthropic/claude-haiku-3-5-20241022": 200_000,
     "anthropic/claude-3-5-sonnet-20241022": 200_000,
     "anthropic/claude-3-opus-20240229": 200_000,
+    "huggingface/fireworks-ai/MiniMaxAI/MiniMax-M2.5": 200_000,
     "huggingface/novita/minimax/minimax-m2.1": 196_608,
     "huggingface/novita/moonshotai/kimi-k2.5": 262_144,
     "huggingface/novita/zai-org/glm-5": 200_000,

agent/core/tools.py

@@ -62,13 +62,7 @@ warnings.filterwarnings(
     "ignore", category=DeprecationWarning, module="aiohttp.connector"
 )
 
-NOT_ALLOWED_TOOL_NAMES = [
-    "hf_jobs",
-    "hf_doc_search",
-    "hf_doc_fetch",
-    "hf_whoami",
-    "paper_search",
-]
+NOT_ALLOWED_TOOL_NAMES = ["hf_jobs", "hf_doc_search", "hf_doc_fetch", "hf_whoami"]
 
 
 def convert_mcp_content_to_string(content: list) -> str:

agent/prompts/system_prompt_v3.yaml

@@ -16,8 +16,8 @@ system_prompt: |
 
   Skip research only for trivial non-code operations.
 
-  When the user asks what is best, optimal, or most effective — start with research papers.
-  Inspect specific resources before using them. Validate dataset format before any training job.
+  For open-ended research tasks (improving model performance, finding the best approach for a task, exploring a field, implementing a paper's method):
+    hf_papers(trending/search) → hf_papers(read_paper) → hf_papers(find_all_resources) → hf_inspect_dataset
 
   # Mistakes you WILL make without research
 
@@ -52,7 +52,6 @@ system_prompt: |
     SFT: "messages", "text", or "prompt"/"completion"
     DPO: "prompt", "chosen", "rejected"
     GRPO: "prompt"
-  All datasets for training should be in conversational ChatML format for HF library compatibility.
 
   # When submitting a training job
 

agent/tools/dataset_tools.py

@@ -10,7 +10,6 @@ from typing import Any, TypedDict
 
 import httpx
 
-from agent.core.session import Session
 from agent.tools.types import ToolResult
 
 BASE_URL = "https://datasets-server.huggingface.co"
@@ -424,17 +423,16 @@ HF_INSPECT_DATASET_TOOL_SPEC = {
 }
 
 
-async def hf_inspect_dataset_handler(
-    arguments: dict[str, Any], session: Session = None
-) -> tuple[str, bool]:
+async def hf_inspect_dataset_handler(arguments: dict[str, Any], session=None) -> tuple[str, bool]:
     """Handler for agent tool router"""
     try:
+        hf_token = session.hf_token if session else None
         result = await inspect_dataset(
             dataset=arguments["dataset"],
             config=arguments.get("config"),
             split=arguments.get("split"),
             sample_rows=min(arguments.get("sample_rows", 3), 10),
-            hf_token=session.hf_token,
+            hf_token=hf_token,
         )
         return result["formatted"], not result.get("isError", False)
     except Exception as e:

backend/dependencies.py

@@ -22,6 +22,9 @@ AUTH_ENABLED = bool(os.environ.get("OAUTH_CLIENT_ID", ""))
 _token_cache: dict[str, tuple[dict[str, Any], float]] = {}
 TOKEN_CACHE_TTL = 300  # 5 minutes
 
+# Org membership cache: key -> expiry_time (only caches positive results)
+_org_member_cache: dict[str, float] = {}
+
 DEV_USER: dict[str, Any] = {
     "user_id": "dev",
     "username": "dev",
@@ -80,6 +83,31 @@ async def _extract_user_from_token(token: str) -> dict[str, Any] | None:
     return None
 
 
+async def check_org_membership(token: str, org_name: str) -> bool:
+    """Check if the token owner belongs to an HF org. Only caches positive results."""
+    now = time.time()
+    key = token + org_name
+    cached = _org_member_cache.get(key)
+    if cached and cached > now:
+        return True
+
+    async with httpx.AsyncClient(timeout=10.0) as client:
+        try:
+            resp = await client.get(
+                f"{OPENID_PROVIDER_URL}/api/whoami-v2",
+                headers={"Authorization": f"Bearer {token}"},
+            )
+            if resp.status_code != 200:
+                return False
+            orgs = {o.get("name") for o in resp.json().get("orgs", [])}
+            if org_name in orgs:
+                _org_member_cache[key] = now + TOKEN_CACHE_TTL
+                return True
+            return False
+        except httpx.HTTPError:
+            return False
+
+
 async def get_current_user(request: Request) -> dict[str, Any]:
     """FastAPI dependency: extract and validate the current user.
 

backend/routes/agent.py

@@ -37,15 +37,15 @@ router = APIRouter(prefix="/api", tags=["agent"])
 
 AVAILABLE_MODELS = [
     {
-        "id": "huggingface/novita/minimax/minimax-m2.1",
-        "label": "MiniMax M2.1",
-        "provider": "huggingface",
+        "id": "anthropic/claude-opus-4-6",
+        "label": "Claude Opus 4.6",
+        "provider": "anthropic",
         "recommended": True,
     },
     {
-        "id": "anthropic/claude-opus-4-5-20251101",
-        "label": "Claude Opus 4.5",
-        "provider": "anthropic",
+        "id": "huggingface/fireworks-ai/MiniMaxAI/MiniMax-M2.5",
+        "label": "MiniMax M2.5",
+        "provider": "huggingface",
         "recommended": True,
     },
     {

backend/routes/auth.py

@@ -142,7 +142,7 @@ async def oauth_callback(
         httponly=True,
         secure=is_production,  # Secure flag only in production (HTTPS)
         samesite="lax",
-        max_age=3600 * 24,  # 24 hours
+        max_age=3600 * 24 * 7,  # 7 days
         path="/",
     )
     return response

backend/session_manager.py

@@ -164,6 +164,17 @@ class SessionManager:
         logger.info(f"Created session {session_id} for user {user_id}")
         return session_id
 
+    @staticmethod
+    async def _cleanup_sandbox(session: Session) -> None:
+        """Delete the sandbox Space if one was created for this session."""
+        sandbox = getattr(session, "sandbox", None)
+        if sandbox and getattr(sandbox, "_owns_space", False):
+            try:
+                logger.info(f"Deleting sandbox {sandbox.space_id}...")
+                await asyncio.to_thread(sandbox.delete)
+            except Exception as e:
+                logger.warning(f"Failed to delete sandbox {sandbox.space_id}: {e}")
+
     async def _run_session(
         self,
         session_id: str,
@@ -218,6 +229,8 @@ class SessionManager:
             except asyncio.CancelledError:
                 pass
 
+            await self._cleanup_sandbox(session)
+
             async with self._lock:
                 if session_id in self.sessions:
                     self.sessions[session_id].is_active = False
@@ -309,6 +322,9 @@ class SessionManager:
 
         ws_manager.clear_buffer(session_id)
 
+        # Clean up sandbox Space before cancelling the task
+        await self._cleanup_sandbox(agent_session.session)
+
         # Cancel the task if running
         if agent_session.task and not agent_session.task.done():
             agent_session.task.cancel()

configs/main_agent_config.json

@@ -1,12 +1,11 @@
 {
-  "model_name": "huggingface/novita/moonshotai/kimi-k2.5",
+  "model_name": "anthropic/claude-opus-4-6",
   "save_sessions": true,
   "session_dataset_repo": "akseljoonas/hf-agent-sessions",
   "yolo_mode": false,
   "confirm_cpu_jobs": true,
   "auto_file_upload": true,
-  "mcpServers": {},
-  "_mcpServers_disabled": {
+  "mcpServers": {
     "hf-mcp-server": {
       "transport": "http",
       "url": "https://huggingface.co/mcp?login"

frontend/src/components/Chat/ActivityStatusBar.tsx

@@ -9,6 +9,7 @@ const shimmer = keyframes`
 `;
 
 const TOOL_LABELS: Record<string, string> = {
+  sandbox_create: 'Creating sandbox, this might take 1-2 minutes',
   hf_jobs: 'Running job',
   hf_repo_files: 'Uploading file',
   hf_repo_git: 'Git operation',

frontend/src/components/Chat/ChatInput.tsx

@@ -21,22 +21,22 @@ const getHfAvatarUrl = (modelId: string) => {
 };
 
 const MODEL_OPTIONS: ModelOption[] = [
-  {
-    id: 'minimax-m2.1',
-    name: 'MiniMax M2.1',
-    description: 'Via Novita',
-    modelPath: 'huggingface/novita/minimax/minimax-m2.1',
-    avatarUrl: getHfAvatarUrl('MiniMaxAI/MiniMax-M2.1'),
-    recommended: true,
-  },
   {
     id: 'claude-opus',
-    name: 'Claude Opus 4.5',
+    name: 'Claude Opus 4.6',
     description: 'Anthropic',
-    modelPath: 'anthropic/claude-opus-4-5-20251101',
+    modelPath: 'anthropic/claude-opus-4-6',
     avatarUrl: 'https://huggingface.co/api/avatars/Anthropic',
     recommended: true,
   },
+  {
+    id: 'minimax-m2.5',
+    name: 'MiniMax M2.5',
+    description: 'Via Fireworks',
+    modelPath: 'huggingface/fireworks-ai/MiniMaxAI/MiniMax-M2.5',
+    avatarUrl: getHfAvatarUrl('MiniMaxAI/MiniMax-M2.5'),
+    recommended: true,
+  },
   {
     id: 'kimi-k2.5',
     name: 'Kimi K2.5',

frontend/src/components/SessionChat.tsx

@@ -94,9 +94,18 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
         store.setActivityStatus({ type: 'tool', toolName: 'running' });
         store.setProcessing(true);
       } else {
-        // No pending approval — reset to idle
-        store.setActivityStatus({ type: 'idle' });
-        store.setProcessing(false);
+        // Check if any tools are still running (non-approval tools like bash, read, etc.)
+        const runningTool = lastAssistant?.parts.find(
+          (p) => p.type === 'dynamic-tool' && (p.state === 'input-available' || p.state === 'input-streaming')
+        );
+        if (runningTool && runningTool.type === 'dynamic-tool') {
+          const desc = (runningTool.input as Record<string, unknown>)?.description as string | undefined;
+          store.setActivityStatus({ type: 'tool', toolName: runningTool.toolName, description: desc });
+          store.setProcessing(true);
+        } else {
+          store.setActivityStatus({ type: 'idle' });
+          store.setProcessing(false);
+        }
       }
     }
     prevActiveRef.current = isActive;

frontend/src/components/WelcomeScreen/WelcomeScreen.tsx

@@ -1,4 +1,4 @@
-import { useState, useCallback } from 'react';
+import { useState, useCallback, useEffect, useRef } from 'react';
 import {
   Box,
   Typography,
@@ -7,6 +7,7 @@ import {
   Alert,
 } from '@mui/material';
 import OpenInNewIcon from '@mui/icons-material/OpenInNew';
+import GroupAddIcon from '@mui/icons-material/GroupAdd';
 import { useSessionStore } from '@/store/sessionStore';
 import { useAgentStore } from '@/store/agentStore';
 import { apiFetch } from '@/utils/api';
@@ -15,29 +16,43 @@ import { isInIframe, triggerLogin } from '@/hooks/useAuth';
 /** HF brand orange */
 const HF_ORANGE = '#FF9D00';
 
+const ORG_JOIN_URL = 'https://huggingface.co/organizations/ml-agent-explorers/share/GzPMJUivoFPlfkvFtIqEouZKSytatKQSZT';
+const ORG_JOINED_KEY = 'hf-agent-org-joined';
+
+function hasJoinedOrg(): boolean {
+  try { return localStorage.getItem(ORG_JOINED_KEY) === '1'; } catch { return false; }
+}
+
+function markOrgJoined(): void {
+  try { localStorage.setItem(ORG_JOINED_KEY, '1'); } catch { /* ignore */ }
+}
+
 export default function WelcomeScreen() {
   const { createSession } = useSessionStore();
   const { setPlan, clearPanel, user } = useAgentStore();
   const [isCreating, setIsCreating] = useState(false);
   const [error, setError] = useState<string | null>(null);
+  const [orgJoined, setOrgJoined] = useState(hasJoinedOrg);
+  const joinLinkOpened = useRef(false);
 
   const inIframe = isInIframe();
   const isAuthenticated = user?.authenticated;
   const isDevUser = user?.username === 'dev';
 
-  const handleStart = useCallback(async () => {
-    if (isCreating) return;
+  // Auto-advance when user returns from the join link
+  useEffect(() => {
+    const handleVisibility = () => {
+      if (document.visibilityState !== 'visible' || !joinLinkOpened.current) return;
+      joinLinkOpened.current = false;
+      markOrgJoined();
+      setOrgJoined(true);
+    };
 
-    // Not authenticated and not dev → need to login
-    if (!isAuthenticated && !isDevUser) {
-      // In iframe: can't redirect (cookies blocked) — user needs to open in new tab
-      // This shouldn't happen because we show a different button in iframe
-      // But just in case:
-      if (inIframe) return;
-      triggerLogin();
-      return;
-    }
+    document.addEventListener('visibilitychange', handleVisibility);
+    return () => document.removeEventListener('visibilitychange', handleVisibility);
+  }, []);
 
+  const tryCreateSession = useCallback(async () => {
     setIsCreating(true);
     setError(null);
 
@@ -65,7 +80,19 @@ export default function WelcomeScreen() {
     } finally {
       setIsCreating(false);
     }
-  }, [isCreating, createSession, setPlan, clearPanel, isAuthenticated, isDevUser, inIframe]);
+  }, [createSession, setPlan, clearPanel]);
+
+  const handleStart = useCallback(async () => {
+    if (isCreating) return;
+
+    if (!isAuthenticated && !isDevUser) {
+      if (inIframe) return;
+      triggerLogin();
+      return;
+    }
+
+    await tryCreateSession();
+  }, [isCreating, isAuthenticated, isDevUser, inIframe, tryCreateSession]);
 
   // Build the direct Space URL for the "open in new tab" link
   const spaceHost = typeof window !== 'undefined'
@@ -74,6 +101,52 @@ export default function WelcomeScreen() {
       : `https://smolagents-ml-agent.hf.space`
     : '';
 
+  // Shared button style
+  const primaryBtnSx = {
+    px: 5,
+    py: 1.5,
+    fontSize: '1rem',
+    fontWeight: 700,
+    textTransform: 'none' as const,
+    borderRadius: '12px',
+    bgcolor: HF_ORANGE,
+    color: '#000',
+    boxShadow: '0 4px 24px rgba(255, 157, 0, 0.3)',
+    textDecoration: 'none',
+    '&:hover': {
+      bgcolor: '#FFB340',
+      boxShadow: '0 6px 32px rgba(255, 157, 0, 0.45)',
+    },
+  };
+
+  // Description block (reused across screens)
+  const description = (
+    <Typography
+      variant="body1"
+      sx={{
+        color: 'var(--muted-text)',
+        maxWidth: 520,
+        mb: 5,
+        lineHeight: 1.8,
+        fontSize: '0.95rem',
+        textAlign: 'center',
+        px: 2,
+        '& strong': { color: 'var(--text)', fontWeight: 600 },
+      }}
+    >
+      A general-purpose AI agent for <strong>machine learning engineering</strong>.
+      It browses <strong>Hugging Face documentation</strong>, manages{' '}
+      <strong>repositories</strong>, launches <strong>training jobs</strong>,
+      and explores <strong>datasets</strong> — all through natural conversation.
+    </Typography>
+  );
+
+  // Which screen to show
+  const needsJoin = inIframe && !orgJoined;
+  const showOpenAgent = inIframe && orgJoined;
+  const showSignin = !inIframe && !isAuthenticated && !isDevUser;
+  const showReady = !inIframe && (isAuthenticated || isDevUser);
+
   return (
     <Box
       sx={{
@@ -109,112 +182,112 @@ export default function WelcomeScreen() {
         HF Agent
       </Typography>
 
-      {/* Description */}
-      <Typography
-        variant="body1"
-        sx={{
-          color: 'var(--muted-text)',
-          maxWidth: 520,
-          mb: 5,
-          lineHeight: 1.8,
-          fontSize: '0.95rem',
-          textAlign: 'center',
-          px: 2,
-          '& strong': { color: 'var(--text)', fontWeight: 600 },
-        }}
-      >
-        A general-purpose AI agent for <strong>machine learning engineering</strong>.
-        It browses <strong>Hugging Face documentation</strong>, manages{' '}
-        <strong>repositories</strong>, launches <strong>training jobs</strong>,
-        and explores <strong>datasets</strong> — all through natural conversation.
-      </Typography>
+      {/* ── Iframe: join org (first visit only) ──────────────────── */}
+      {needsJoin && (
+        <>
+          <Typography
+            variant="body1"
+            sx={{
+              color: 'var(--muted-text)',
+              maxWidth: 480,
+              mb: 4,
+              lineHeight: 1.8,
+              fontSize: '0.95rem',
+              textAlign: 'center',
+              px: 2,
+              '& strong': { color: 'var(--text)', fontWeight: 600 },
+            }}
+          >
+            Under the hood, this agent uses GPUs, inference APIs, and other paid Hub goodies — but we made them all free for you. Just join <strong>ML Agent Explorers</strong> to get started!
+          </Typography>
 
-      {/* Action button — depends on context */}
-      {inIframe && !isAuthenticated && !isDevUser ? (
-        // In iframe + not logged in → link to open Space directly
-        <Button
-          variant="contained"
-          size="large"
-          component="a"
-          href={spaceHost}
-          target="_blank"
-          rel="noopener noreferrer"
-          endIcon={<OpenInNewIcon />}
-          sx={{
-            px: 5,
-            py: 1.5,
-            fontSize: '1rem',
-            fontWeight: 700,
-            textTransform: 'none',
-            borderRadius: '12px',
-            bgcolor: HF_ORANGE,
-            color: '#000',
-            boxShadow: '0 4px 24px rgba(255, 157, 0, 0.3)',
-            textDecoration: 'none',
-            '&:hover': {
-              bgcolor: '#FFB340',
-              boxShadow: '0 6px 32px rgba(255, 157, 0, 0.45)',
-            },
-          }}
-        >
-          Open HF Agent
-        </Button>
-      ) : !isAuthenticated && !isDevUser ? (
-        // Direct access + not logged in → sign in button
-        <Button
-          variant="contained"
-          size="large"
-          onClick={() => triggerLogin()}
-          sx={{
-            px: 5,
-            py: 1.5,
-            fontSize: '1rem',
-            fontWeight: 700,
-            textTransform: 'none',
-            borderRadius: '12px',
-            bgcolor: HF_ORANGE,
-            color: '#000',
-            boxShadow: '0 4px 24px rgba(255, 157, 0, 0.3)',
-            '&:hover': {
-              bgcolor: '#FFB340',
-              boxShadow: '0 6px 32px rgba(255, 157, 0, 0.45)',
-            },
-          }}
-        >
-          Sign in with Hugging Face
-        </Button>
-      ) : (
-        // Authenticated or dev → start session
-        <Button
-          variant="contained"
-          size="large"
-          onClick={handleStart}
-          disabled={isCreating}
-          startIcon={
-            isCreating ? <CircularProgress size={20} color="inherit" /> : null
-          }
-          sx={{
-            px: 5,
-            py: 1.5,
-            fontSize: '1rem',
-            fontWeight: 700,
-            textTransform: 'none',
-            borderRadius: '12px',
-            bgcolor: HF_ORANGE,
-            color: '#000',
-            boxShadow: '0 4px 24px rgba(255, 157, 0, 0.3)',
-            '&:hover': {
-              bgcolor: '#FFB340',
-              boxShadow: '0 6px 32px rgba(255, 157, 0, 0.45)',
-            },
-            '&.Mui-disabled': {
-              bgcolor: 'rgba(255, 157, 0, 0.35)',
-              color: 'rgba(0,0,0,0.45)',
-            },
-          }}
-        >
-          {isCreating ? 'Initializing...' : 'Start Session'}
-        </Button>
+          <Button
+            variant="contained"
+            size="large"
+            component="a"
+            href={ORG_JOIN_URL}
+            target="_blank"
+            rel="noopener noreferrer"
+            onClick={() => { joinLinkOpened.current = true; }}
+            startIcon={<GroupAddIcon />}
+            sx={primaryBtnSx}
+          >
+            Join ML Agent Explorers
+          </Button>
+        </>
+      )}
+
+      {/* ── Iframe: already joined → open Space ──────────────────── */}
+      {showOpenAgent && (
+        <>
+          {description}
+          <Button
+            variant="contained"
+            size="large"
+            component="a"
+            href={spaceHost}
+            target="_blank"
+            rel="noopener noreferrer"
+            endIcon={<OpenInNewIcon />}
+            sx={primaryBtnSx}
+          >
+            Open HF Agent
+          </Button>
+        </>
+      )}
+
+      {/* ── Direct: not logged in → sign in ───────────────���──────── */}
+      {showSignin && (
+        <>
+          {description}
+          <Button
+            variant="contained"
+            size="large"
+            onClick={() => triggerLogin()}
+            sx={primaryBtnSx}
+          >
+            Sign in with Hugging Face
+          </Button>
+
+          <Typography
+            variant="caption"
+            sx={{
+              mt: 2.5,
+              color: 'var(--muted-text)',
+              fontSize: '0.78rem',
+              textAlign: 'center',
+              maxWidth: 360,
+              lineHeight: 1.6,
+            }}
+          >
+            Make sure to enable access to the <strong>ml-agent-explorers</strong> org when prompted.
+          </Typography>
+        </>
+      )}
+
+      {/* ── Direct: authenticated → start session ────────────────── */}
+      {showReady && (
+        <>
+          {description}
+          <Button
+            variant="contained"
+            size="large"
+            onClick={handleStart}
+            disabled={isCreating}
+            startIcon={
+              isCreating ? <CircularProgress size={20} color="inherit" /> : null
+            }
+            sx={{
+              ...primaryBtnSx,
+              '&.Mui-disabled': {
+                bgcolor: 'rgba(255, 157, 0, 0.35)',
+                color: 'rgba(0,0,0,0.45)',
+              },
+            }}
+          >
+            {isCreating ? 'Initializing...' : 'Start Session'}
+          </Button>
+        </>
       )}
 
       {/* Error */}