Remember the chosen jobs namespace and clean up the picker dialog (#162)

When a user has paid orgs but no Pro subscription, hf_jobs needs an
explicit namespace. The picker fired on every tool call — even after
the user already chose one. Persist the selection in localStorage and
auto-attach it to subsequent hf_jobs approvals; if the saved namespace
becomes invalid (e.g. left the org), backend 400 surfaces a typed error
that clears the preference and reopens the picker.

UI is rewritten to match ClaudeCapDialog: tighter copy, single dropdown
without the redundant "Eligible namespaces" panel, "Continue" + "Skip
this tool call" actions.

backend/routes/agent.py

@@ -195,6 +195,9 @@ async def _enforce_jobs_access_for_approvals(
                         "The selected jobs namespace is not one of your eligible paid organizations. "
                         f"Allowed namespaces: {', '.join(access.paid_org_names)}"
                     ),
+                    "plan": user.get("plan", "free"),
+                    "tool_call_ids": invalid_namespace,
+                    "eligible_namespaces": access.paid_org_names,
                 },
             )
         missing_namespace = [

frontend/src/components/JobsUpgradeDialog.tsx

@@ -8,7 +8,6 @@ import {
   DialogContentText,
   DialogTitle,
   FormControl,
-  InputLabel,
   MenuItem,
   Select,
   Typography,
@@ -37,13 +36,20 @@ export default function JobsUpgradeDialog({
   onClose,
   onContinueWithNamespace,
 }: JobsUpgradeDialogProps) {
-  const [selectedNamespace, setSelectedNamespace] = useState('');
+  const [selectedNamespace, setSelectedNamespace] = useState(() => eligibleNamespaces[0] || '');
 
   useEffect(() => {
     if (!open) return;
     setSelectedNamespace(eligibleNamespaces[0] || '');
   }, [open, eligibleNamespaces]);
 
+  const isNamespace = mode === 'namespace';
+  const title = isNamespace ? 'Run jobs as' : 'Jobs need Pro or a paid org';
+
+  const body = isNamespace
+    ? "Pick which paid organization should pay for and own this job. We'll use the same one for the rest of this browser."
+    : message;
+
   return (
     <Dialog
       open={open}
@@ -57,7 +63,7 @@ export default function JobsUpgradeDialog({
           border: '1px solid var(--border)',
           borderRadius: 'var(--radius-md)',
           boxShadow: 'var(--shadow-1)',
-          maxWidth: 500,
+          maxWidth: 460,
           mx: 2,
         },
       }}
@@ -65,72 +71,75 @@ export default function JobsUpgradeDialog({
       <DialogTitle
         sx={{ color: 'var(--text)', fontWeight: 700, fontSize: '1rem', pt: 2.5, pb: 0, px: 3 }}
       >
-        {mode === 'namespace' ? 'Choose the org for this job' : 'Jobs need Pro or a paid org'}
+        {title}
       </DialogTitle>
       <DialogContent sx={{ px: 3, pt: 1.25, pb: 0 }}>
         <DialogContentText
           sx={{ color: 'var(--muted-text)', fontSize: '0.85rem', lineHeight: 1.6 }}
         >
-          {message}
+          {body}
         </DialogContentText>
-        {eligibleNamespaces.length > 0 && (
-          <Box
-            sx={{
-              mt: 2,
-              p: 1.5,
-              borderRadius: '8px',
-              bgcolor: 'var(--accent-yellow-weak)',
-              border: '1px solid var(--border)',
-            }}
-          >
-            <Typography
-              variant="caption"
+
+        {isNamespace ? (
+          <FormControl fullWidth size="small" sx={{ mt: 2 }}>
+            <Select
+              value={selectedNamespace}
+              displayEmpty
+              onChange={(e) => setSelectedNamespace(String(e.target.value))}
               sx={{
-                display: 'block',
-                fontWeight: 700,
+                bgcolor: 'var(--composer-bg)',
                 color: 'var(--text)',
-                fontSize: '0.78rem',
-                mb: 1,
-                letterSpacing: '0.02em',
+                fontSize: '0.88rem',
+                fontWeight: 600,
+                '& .MuiOutlinedInput-notchedOutline': { borderColor: 'var(--border)' },
+                '&:hover .MuiOutlinedInput-notchedOutline': { borderColor: 'var(--border)' },
+                '&.Mui-focused .MuiOutlinedInput-notchedOutline': {
+                  borderColor: 'var(--accent-yellow)',
+                  borderWidth: 1,
+                },
+                '& .MuiSelect-icon': { color: 'var(--muted-text)' },
+              }}
+              MenuProps={{
+                PaperProps: {
+                  sx: {
+                    bgcolor: 'var(--panel)',
+                    border: '1px solid var(--border)',
+                    borderRadius: '8px',
+                    mt: 0.5,
+                  },
+                },
               }}
             >
-              Eligible namespaces
-            </Typography>
-            {mode === 'namespace' ? (
-              <FormControl fullWidth size="small">
-                <InputLabel id="jobs-namespace-label">Organization</InputLabel>
-                <Select
-                  labelId="jobs-namespace-label"
-                  value={selectedNamespace}
-                  label="Organization"
-                  onChange={(e) => setSelectedNamespace(String(e.target.value))}
+              {eligibleNamespaces.map((namespace) => (
+                <MenuItem
+                  key={namespace}
+                  value={namespace}
+                  sx={{
+                    fontSize: '0.88rem',
+                    color: 'var(--text)',
+                    '&.Mui-selected': { bgcolor: 'rgba(255,255,255,0.05)' },
+                  }}
                 >
-                  {eligibleNamespaces.map((namespace) => (
-                    <MenuItem key={namespace} value={namespace}>
-                      {namespace}
-                    </MenuItem>
-                  ))}
-                </Select>
-              </FormControl>
-            ) : (
+                  {namespace}
+                </MenuItem>
+              ))}
+            </Select>
+          </FormControl>
+        ) : (
+          eligibleNamespaces.length > 0 && (
+            <Box sx={{ mt: 1.5 }}>
               <Typography
                 variant="caption"
-                sx={{ display: 'block', color: 'var(--muted-text)', fontSize: '0.78rem', lineHeight: 1.55 }}
+                sx={{ color: 'var(--muted-text)', fontSize: '0.78rem', lineHeight: 1.55 }}
               >
-                {eligibleNamespaces.join(', ')}
+                Eligible namespaces: {eligibleNamespaces.join(', ')}
               </Typography>
-            )}
-          </Box>
+            </Box>
+          )
         )}
-        <Typography
-          variant="caption"
-          sx={{ display: 'block', mt: 2, color: 'var(--muted-text)', fontSize: '0.78rem', lineHeight: 1.55 }}
-        >
-          If you decline, the agent will have to find another way forward without `hf_jobs`.
-        </Typography>
       </DialogContent>
-      <DialogActions sx={{ px: 3, pb: 2.5, pt: 2, gap: 1 }}>
-        {mode === 'namespace' ? (
+      <DialogActions sx={{ px: 3, pb: 2.5, pt: 2.5, gap: 1 }}>
+        {isNamespace ? (
           <Button
             onClick={() => onContinueWithNamespace(selectedNamespace)}
             disabled={!selectedNamespace}
@@ -147,7 +156,7 @@ export default function JobsUpgradeDialog({
               '&:hover': { bgcolor: '#FFB340', boxShadow: 'none' },
             }}
           >
-            Run under selected org
+            Continue
           </Button>
         ) : (
           <Button
@@ -183,7 +192,7 @@ export default function JobsUpgradeDialog({
             '&:hover': { bgcolor: 'var(--hover-bg)' },
           }}
         >
-          Decline tool call
+          {isNamespace ? 'Skip this tool call' : 'Decline tool call'}
         </Button>
       </DialogActions>
     </Dialog>

frontend/src/hooks/useAgentChat.ts

@@ -447,6 +447,33 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         }
         return;
       }
+      if (error.message === 'HF_JOBS_INVALID_NAMESPACE') {
+        // Saved preference is no longer one of the user's eligible namespaces
+        // (e.g. they left the org). Clear it and reopen the picker.
+        const typed = error as Error & {
+          detail?: Record<string, unknown>;
+          approvals?: Array<{
+            tool_call_id: string;
+            approved: boolean;
+            feedback?: string | null;
+            edited_script?: string | null;
+            namespace?: string | null;
+          }>;
+        };
+        useAgentStore.getState().setPreferredJobsNamespace(null);
+        void hydrateFromBackend();
+        if (isActiveRef.current) {
+          useAgentStore.getState().setJobsUpgradeRequired({
+            approvals: typed.approvals || [],
+            toolCallIds: (typed.detail?.tool_call_ids as string[]) || [],
+            message: String(typed.detail?.message || 'Pick a different organization for this job run.'),
+            eligibleNamespaces: (typed.detail?.eligible_namespaces as string[]) || [],
+            plan: ((typed.detail?.plan as 'free' | 'pro' | 'org') || 'free'),
+            mode: 'namespace',
+          });
+        }
+        return;
+      }
       logger.error('useChat error:', error);
       if (isActiveRef.current) {
         useAgentStore.getState().setError(error.message);
@@ -830,6 +857,9 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
         : approval.namespace,
     }));
 
+    // Remember this choice so the picker doesn't reappear for every
+    // subsequent hf_jobs call.
+    useAgentStore.getState().setPreferredJobsNamespace(namespace);
     useAgentStore.getState().setJobsUpgradeRequired(null);
     return approveTools(approvals);
   }, [approveTools]);

frontend/src/lib/sse-chat-transport.ts

@@ -325,7 +325,14 @@ export class SSEChatTransport implements ChatTransport<UIMessage> {
         const approved = p.approval?.approved ?? true;
         // Get edited script from agentStore if available
         const editedScript = useAgentStore.getState().getEditedScript(p.toolCallId);
-        const namespace = useAgentStore.getState().getApprovalNamespace(p.toolCallId);
+        const explicitNamespace = useAgentStore.getState().getApprovalNamespace(p.toolCallId);
+        // Fall back to the user's persisted choice so we don't re-prompt
+        // every hf_jobs call.  Backend will 400 if the saved namespace is
+        // no longer valid; the error handler clears the preference and
+        // reopens the picker.
+        const preferred = useAgentStore.getState().preferredJobsNamespace;
+        const namespace = explicitNamespace
+          ?? (approved && p.toolName === 'hf_jobs' ? preferred ?? null : null);
         return {
           tool_call_id: p.toolCallId,
           approved,
@@ -393,6 +400,20 @@ export class SSEChatTransport implements ChatTransport<UIMessage> {
         throw err;
       }
     }
+    if (response.status === 400) {
+      const payload = await response.json().catch(() => null);
+      if (payload?.detail?.error === 'hf_jobs_invalid_namespace') {
+        // Stored namespace is no longer eligible — surface so the UI can
+        // clear the saved preference and reopen the picker.
+        const err = new Error('HF_JOBS_INVALID_NAMESPACE') as Error & {
+          detail?: Record<string, unknown>;
+          approvals?: Array<Record<string, unknown>>;
+        };
+        err.detail = payload.detail as Record<string, unknown>;
+        err.approvals = (body.approvals as Array<Record<string, unknown>> | undefined) || [];
+        throw err;
+      }
+    }
     if (!response.ok) {
       const errorText = await response.text().catch(() => 'Request failed');
       throw new Error(`Chat request failed: ${response.status} ${errorText}`);

frontend/src/store/agentStore.ts

@@ -141,6 +141,10 @@ interface AgentStore {
   // Namespace overrides chosen for hf_jobs approvals (tool_call_id -> namespace)
   approvalNamespaces: Record<string, string>;
 
+  // Persisted preferred namespace for hf_jobs (auto-applied to future approvals
+  // so the user only picks once)
+  preferredJobsNamespace: string | null;
+
   // Job URLs (tool_call_id -> job URL) for HF jobs
   jobUrls: Record<string, string>;
 
@@ -199,6 +203,8 @@ interface AgentStore {
   getApprovalNamespace: (toolCallId: string) => string | undefined;
   clearApprovalNamespaces: () => void;
 
+  setPreferredJobsNamespace: (namespace: string | null) => void;
+
   setJobUrl: (toolCallId: string, jobUrl: string) => void;
   getJobUrl: (toolCallId: string) => string | undefined;
 
@@ -292,6 +298,28 @@ function saveTrackioDashboards(dashboards: Record<string, { spaceId: string; pro
   }
 }
 
+const PREFERRED_JOBS_NAMESPACE_KEY = 'hf-agent-preferred-jobs-namespace';
+
+function loadPreferredJobsNamespace(): string | null {
+  try {
+    return localStorage.getItem(PREFERRED_JOBS_NAMESPACE_KEY);
+  } catch {
+    return null;
+  }
+}
+
+function savePreferredJobsNamespace(namespace: string | null): void {
+  try {
+    if (namespace) {
+      localStorage.setItem(PREFERRED_JOBS_NAMESPACE_KEY, namespace);
+    } else {
+      localStorage.removeItem(PREFERRED_JOBS_NAMESPACE_KEY);
+    }
+  } catch (e) {
+    console.warn('Failed to persist preferred jobs namespace:', e);
+  }
+}
+
 export const useAgentStore = create<AgentStore>()((set, get) => ({
   sessionStates: {},
   activeSessionId: null,
@@ -313,6 +341,7 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
 
   editedScripts: {},
   approvalNamespaces: {},
+  preferredJobsNamespace: loadPreferredJobsNamespace(),
   jobUrls: {},
   jobStatuses: {},
   trackioDashboards: loadTrackioDashboards(),
@@ -494,6 +523,11 @@ export const useAgentStore = create<AgentStore>()((set, get) => ({
 
   clearApprovalNamespaces: () => set({ approvalNamespaces: {} }),
 
+  setPreferredJobsNamespace: (namespace) => {
+    savePreferredJobsNamespace(namespace);
+    set({ preferredJobsNamespace: namespace });
+  },
+
   // ── Job URLs ────────────────────────────────────────────────────────
 
   setJobUrl: (toolCallId, jobUrl) => {