Upload 102 files

ui/src/app/api/auth/hf/exchange/route.ts

@@ -0,0 +1,80 @@
+import { NextRequest, NextResponse } from 'next/server';
+
+const TOKEN_ENDPOINT = 'https://huggingface.co/oauth/token';
+const USERINFO_ENDPOINT = 'https://huggingface.co/oauth/userinfo';
+const STATE_COOKIE = 'hf_oauth_state';
+
+export async function POST(request: NextRequest) {
+  const clientId = process.env.HF_OAUTH_CLIENT_ID || process.env.NEXT_PUBLIC_HF_OAUTH_CLIENT_ID;
+  const clientSecret = process.env.HF_OAUTH_CLIENT_SECRET;
+
+  if (!clientId || !clientSecret) {
+    return NextResponse.json({ error: 'OAuth application is not configured' }, { status: 500 });
+  }
+
+  const { code, state } = await request.json().catch(() => ({}));
+
+  if (!code) {
+    return NextResponse.json({ error: 'Authorization code is missing' }, { status: 400 });
+  }
+
+  const storedState = request.cookies.get(STATE_COOKIE)?.value;
+  if (!storedState || state !== storedState) {
+    const response = NextResponse.json({ error: 'Invalid or expired OAuth state' }, { status: 400 });
+    response.cookies.delete(STATE_COOKIE);
+    return response;
+  }
+
+  const origin = request.nextUrl.origin;
+  const redirectUri =
+    process.env.HF_OAUTH_REDIRECT_URI ||
+    process.env.NEXT_PUBLIC_HF_OAUTH_REDIRECT_URI ||
+    `${origin}/auth/hf/callback`;
+
+  try {
+    const tokenResponse = await fetch(TOKEN_ENDPOINT, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
+      body: new URLSearchParams({
+        grant_type: 'authorization_code',
+        code,
+        redirect_uri: redirectUri,
+        client_id: clientId,
+        client_secret: clientSecret,
+      }),
+    });
+
+    if (!tokenResponse.ok) {
+      const errorPayload = await tokenResponse.json().catch(() => ({}));
+      throw new Error(errorPayload?.error_description || 'Failed to exchange code for token');
+    }
+
+    const tokenData = await tokenResponse.json();
+    const accessToken = tokenData?.access_token;
+    if (!accessToken) {
+      throw new Error('Access token missing in response');
+    }
+
+    const userResponse = await fetch(USERINFO_ENDPOINT, {
+      headers: { Authorization: `Bearer ${accessToken}` },
+    });
+
+    if (!userResponse.ok) {
+      throw new Error('Failed to fetch user info');
+    }
+
+    const profile = await userResponse.json();
+    const namespace = profile?.preferred_username || profile?.name || 'user';
+
+    const response = NextResponse.json({
+      token: accessToken,
+      namespace,
+    });
+    response.cookies.delete(STATE_COOKIE);
+    return response;
+  } catch (error: any) {
+    const response = NextResponse.json({ error: error?.message || 'OAuth flow failed' }, { status: 500 });
+    response.cookies.delete(STATE_COOKIE);
+    return response;
+  }
+}

ui/src/app/api/auth/hf/login/route.ts

@@ -12,7 +12,10 @@ export async function GET(request: NextRequest) {
 
   const state = randomUUID();
   const origin = request.nextUrl.origin;
-  const redirectUri = process.env.HF_OAUTH_REDIRECT_URI || process.env.NEXT_PUBLIC_HF_OAUTH_REDIRECT_URI || `${origin}/api/auth/hf/callback`;
+  const redirectUri =
+    process.env.HF_OAUTH_REDIRECT_URI ||
+    process.env.NEXT_PUBLIC_HF_OAUTH_REDIRECT_URI ||
+    `${origin}/auth/hf/callback`;
 
   const authorizeUrl = new URL(HF_AUTHORIZE_URL);
   authorizeUrl.searchParams.set('response_type', 'code');

ui/src/app/auth/hf/callback/page.tsx

@@ -0,0 +1,54 @@
+'use client';
+
+import { useEffect, useState } from 'react';
+import { useRouter, useSearchParams } from 'next/navigation';
+import { useAuth } from '@/contexts/AuthContext';
+import Loading from '@/components/Loading';
+
+export default function HFOAuthCallbackPage() {
+  const router = useRouter();
+  const searchParams = useSearchParams();
+  const { exchangeCodeForToken, error } = useAuth();
+  const [localError, setLocalError] = useState<string | null>(null);
+
+  useEffect(() => {
+    const code = searchParams.get('code');
+    const state = searchParams.get('state');
+    const errorParam = searchParams.get('error');
+
+    const handleAuth = async () => {
+      if (errorParam) {
+        setLocalError(errorParam);
+        router.replace('/settings');
+        return;
+      }
+
+      if (code && state) {
+        const success = await exchangeCodeForToken(code, state);
+        if (success) {
+          router.replace('/dashboard');
+        } else {
+          setLocalError('Authentication failed. Please try again.');
+          router.replace('/settings');
+        }
+      } else {
+        setLocalError('Invalid authentication callback.');
+        router.replace('/settings');
+      }
+    };
+
+    handleAuth();
+  }, [exchangeCodeForToken, router, searchParams]);
+
+  return (
+    <div className="min-h-screen flex flex-col items-center justify-center text-gray-200 bg-gray-950 gap-4">
+      <Loading />
+      <p className="text-lg">Authenticating with Hugging Face…</p>
+      {(localError || error) && (
+        <p className="text-sm text-red-400 max-w-md text-center">
+          {localError || error}
+        </p>
+      )}
+    </div>
+  );
+}

ui/src/app/jobs/new/page.tsx

@@ -109,7 +109,13 @@ export default function TrainingForm() {
     return () => {
       isMounted = false;
     };
-  }, [datasets, settings, isSettingsLoaded, datasetFetchStatus]);
+  }, [datasets, settings, isSettingsLoaded, datasetFetchStatus, isAuthenticated]);
+
+  useEffect(() => {
+    if (!isAuthenticated) {
+      setDatasetOptions([]);
+    }
+  }, [isAuthenticated]);
 
   useEffect(() => {
     if (runId) {
@@ -299,8 +305,3 @@ export default function TrainingForm() {
     </>
   );
 }
-  useEffect(() => {
-    if (!isAuthenticated) {
-      setDatasetOptions([]);
-    }
-  }, [isAuthenticated]);

ui/src/contexts/AuthContext.tsx

@@ -22,6 +22,7 @@ interface AuthContextValue {
   error: string | null;
   oauthAvailable: boolean;
   loginWithOAuth: () => void;
+  exchangeCodeForToken: (code: string, state: string) => Promise<boolean>;
   setManualToken: (token: string) => Promise<void>;
   logout: () => void;
 }
@@ -36,6 +37,7 @@ const defaultValue: AuthContextValue = {
   error: null,
   oauthAvailable: Boolean(oauthClientId),
   loginWithOAuth: () => {},
+  exchangeCodeForToken: async () => false,
   setManualToken: async () => {},
   logout: () => {},
 };
@@ -191,6 +193,46 @@ export function AuthProvider({ children }: { children: React.ReactNode }) {
     [applyAuthState],
   );
 
+  const exchangeCodeForToken = useCallback(
+    async (code: string, state: string) => {
+      if (!code || !state) {
+        setError('Invalid authorization response.');
+        setStatus('error');
+        return false;
+      }
+      setStatus('checking');
+      setError(null);
+      try {
+        const res = await fetch('/api/auth/hf/exchange', {
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+          },
+          credentials: 'same-origin',
+          body: JSON.stringify({ code, state }),
+        });
+
+        if (!res.ok) {
+          const data = await res.json().catch(() => ({}));
+          throw new Error(data?.error || 'Failed to exchange authorization code');
+        }
+
+        const data = await res.json();
+        await applyAuthState({
+          token: data.token,
+          namespace: data.namespace || 'user',
+          method: 'oauth',
+        });
+        return true;
+      } catch (err: any) {
+        setError(err?.message || 'Failed to authenticate with Hugging Face');
+        setStatus('error');
+        return false;
+      }
+    },
+    [applyAuthState],
+  );
+
   const loginWithOAuth = useCallback(() => {
     if (typeof window === 'undefined') {
       return;
@@ -203,54 +245,13 @@ export function AuthProvider({ children }: { children: React.ReactNode }) {
     setStatus('checking');
     setError(null);
 
-    const width = 540;
-    const height = 720;
-    const left = window.screenX + (window.outerWidth - width) / 2;
-    const top = window.screenY + (window.outerHeight - height) / 2;
-
-    window.open(
-      '/api/auth/hf/login',
-      'hf-oauth-window',
-      `width=${width},height=${height},left=${left},top=${top},resizable,scrollbars=yes,status=1`,
-    );
+    window.location.href = '/api/auth/hf/login';
   }, []);
 
   const logout = useCallback(() => {
     clearAuthState();
   }, [clearAuthState]);
 
-  // Listen for OAuth completion messages
-  useEffect(() => {
-    if (typeof window === 'undefined') {
-      return;
-    }
-
-    const handler = async (event: MessageEvent) => {
-      if (event.origin !== window.location.origin) {
-        return;
-      }
-
-      const { type, payload } = event.data || {};
-
-      if (type === 'HF_OAUTH_SUCCESS') {
-        await applyAuthState({
-          token: payload?.token,
-          namespace: payload?.namespace || 'user',
-          method: 'oauth',
-        });
-        return;
-      }
-
-      if (type === 'HF_OAUTH_ERROR') {
-        setStatus('error');
-        setError(payload?.message || 'OAuth flow failed');
-      }
-    };
-
-    window.addEventListener('message', handler);
-    return () => window.removeEventListener('message', handler);
-  }, [applyAuthState]);
-
   const value = useMemo<AuthContextValue>(
     () => ({
       status,
@@ -260,10 +261,11 @@ export function AuthProvider({ children }: { children: React.ReactNode }) {
       error,
       oauthAvailable,
       loginWithOAuth,
+      exchangeCodeForToken,
       setManualToken,
       logout,
     }),
-    [status, token, namespace, method, error, oauthAvailable, loginWithOAuth, setManualToken, logout],
+    [status, token, namespace, method, error, oauthAvailable, loginWithOAuth, exchangeCodeForToken, setManualToken, logout],
   );
 
   return <AuthContext.Provider value={value}>{children}</AuthContext.Provider>;