From c038cabaf6a4de6c400537a9a90a951ece1c3a14 Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Thu, 31 Jul 2025 18:59:12 +0200
Subject: [PATCH] Implement authentication flow with API key support and
 loading states

---
 webui/src/App.tsx                       |  29 ++++-
 webui/src/components/Header.tsx         |  27 +++-
 webui/src/components/LoginDialog.tsx    | 157 +++++++++++++++++++++++
 webui/src/contexts/AuthContext.tsx      | 162 ++++++++++++++++++++++++
 webui/src/contexts/InstancesContext.tsx |  26 +++-
 webui/src/lib/api.ts                    |  21 ++-
 webui/src/main.tsx                      |   9 +-
 7 files changed, 414 insertions(+), 17 deletions(-)
 create mode 100644 webui/src/components/LoginDialog.tsx
 create mode 100644 webui/src/contexts/AuthContext.tsx

diff --git a/webui/src/App.tsx b/webui/src/App.tsx
index 3822554..6ae4e78 100644
--- a/webui/src/App.tsx
+++ b/webui/src/App.tsx
@@ -2,11 +2,14 @@ import { useState } from "react";
 import Header from "@/components/Header";
 import InstanceList from "@/components/InstanceList";
 import InstanceModal from "@/components/InstanceModal";
+import LoginDialog from "@/components/LoginDialog";
+import SystemInfoModal from "./components/SystemInfoModal";
 import { type CreateInstanceOptions, type Instance } from "@/types/instance";
 import { useInstances } from "@/contexts/InstancesContext";
-import SystemInfoModal from "./components/SystemInfoModal";
+import { useAuth } from "@/contexts/AuthContext";
 
 function App() {
+  const { isAuthenticated, isLoading: authLoading } = useAuth();
   const [isInstanceModalOpen, setIsInstanceModalOpen] = useState(false);
   const [isSystemInfoModalOpen, setIsSystemInfoModalOpen] = useState(false);
   const [editingInstance, setEditingInstance] = useState<Instance | undefined>(
@@ -36,6 +39,28 @@ function App() {
     setIsSystemInfoModalOpen(true);
   };
 
+  // Show loading spinner while checking auth
+  if (authLoading) {
+    return (
+      <div className="min-h-screen bg-gray-50 flex items-center justify-center">
+        <div className="text-center">
+          <div className="animate-spin rounded-full h-8 w-8 border-b-2 border-blue-600 mx-auto mb-4"></div>
+          <p className="text-gray-600">Loading...</p>
+        </div>
+      </div>
+    );
+  }
+
+  // Show login dialog if not authenticated
+  if (!isAuthenticated) {
+    return (
+      <div className="min-h-screen bg-gray-50">
+        <LoginDialog open={true} />
+      </div>
+    );
+  }
+
+  // Show main app if authenticated
   return (
     <div className="min-h-screen bg-gray-50">
       <Header onCreateInstance={handleCreateInstance} onShowSystemInfo={handleShowSystemInfo} />
@@ -58,4 +83,4 @@ function App() {
   );
 }
 
-export default App;
+export default App;
\ No newline at end of file
diff --git a/webui/src/components/Header.tsx b/webui/src/components/Header.tsx
index 544677c..ed272ed 100644
--- a/webui/src/components/Header.tsx
+++ b/webui/src/components/Header.tsx
@@ -1,5 +1,6 @@
 import { Button } from "@/components/ui/button";
-import { HelpCircle } from "lucide-react";
+import { HelpCircle, LogOut } from "lucide-react";
+import { useAuth } from "@/contexts/AuthContext";
 
 interface HeaderProps {
   onCreateInstance: () => void;
@@ -7,6 +8,14 @@ interface HeaderProps {
 }
 
 function Header({ onCreateInstance, onShowSystemInfo }: HeaderProps) {
+  const { logout } = useAuth();
+
+  const handleLogout = () => {
+    if (confirm("Are you sure you want to logout?")) {
+      logout();
+    }
+  };
+
   return (
     <header className="bg-white border-b border-gray-200">
       <div className="container mx-auto max-w-4xl px-4 py-4">
@@ -16,7 +25,9 @@ function Header({ onCreateInstance, onShowSystemInfo }: HeaderProps) {
           </h1>
 
           <div className="flex items-center gap-2">
-            <Button onClick={onCreateInstance} data-testid="create-instance-button">Create Instance</Button>
+            <Button onClick={onCreateInstance} data-testid="create-instance-button">
+              Create Instance
+            </Button>
 
             <Button
               variant="outline"
@@ -27,6 +38,16 @@ function Header({ onCreateInstance, onShowSystemInfo }: HeaderProps) {
             >
               <HelpCircle className="h-4 w-4" />
             </Button>
+
+            <Button
+              variant="outline"
+              size="icon"
+              onClick={handleLogout}
+              data-testid="logout-button"
+              title="Logout"
+            >
+              <LogOut className="h-4 w-4" />
+            </Button>
           </div>
         </div>
       </div>
@@ -34,4 +55,4 @@ function Header({ onCreateInstance, onShowSystemInfo }: HeaderProps) {
   );
 }
 
-export default Header;
+export default Header;
\ No newline at end of file
diff --git a/webui/src/components/LoginDialog.tsx b/webui/src/components/LoginDialog.tsx
new file mode 100644
index 0000000..3f29948
--- /dev/null
+++ b/webui/src/components/LoginDialog.tsx
@@ -0,0 +1,157 @@
+import React, { useState, useEffect } from 'react'
+import { Button } from '@/components/ui/button'
+import { Input } from '@/components/ui/input'
+import { Label } from '@/components/ui/label'
+import {
+  Dialog,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+} from '@/components/ui/dialog'
+import { AlertCircle, Key, Loader2 } from 'lucide-react'
+import { useAuth } from '@/contexts/AuthContext'
+
+interface LoginDialogProps {
+  open: boolean
+  onOpenChange?: (open: boolean) => void
+}
+
+const LoginDialog: React.FC<LoginDialogProps> = ({
+  open,
+  onOpenChange,
+}) => {
+  const { login, isLoading, error, clearError } = useAuth()
+  const [apiKey, setApiKey] = useState('')
+  const [localLoading, setLocalLoading] = useState(false)
+
+  // Clear form and errors when dialog opens/closes
+  useEffect(() => {
+    if (open) {
+      setApiKey('')
+      clearError()
+    }
+  }, [open, clearError])
+
+  // Clear error when user starts typing
+  useEffect(() => {
+    if (error && apiKey) {
+      clearError()
+    }
+  }, [apiKey, error, clearError])
+
+  const handleSubmit = async (e: React.FormEvent) => {
+    e.preventDefault()
+    
+    if (!apiKey.trim()) {
+      return
+    }
+
+    setLocalLoading(true)
+    
+    try {
+      await login(apiKey.trim())
+      // Login successful - dialog will close automatically when auth state changes
+      setApiKey('')
+    } catch (err) {
+      // Error is handled by the AuthContext
+      console.error('Login failed:', err)
+    } finally {
+      setLocalLoading(false)
+    }
+  }
+
+  const handleCancel = () => {
+    setApiKey('')
+    clearError()
+    onOpenChange?.(false)
+  }
+
+  const handleKeyDown = (e: React.KeyboardEvent<HTMLInputElement>) => {
+    if (e.key === 'Enter' && !isSubmitDisabled) {
+      // Create a synthetic FormEvent to satisfy handleSubmit's type
+      const syntheticEvent = {
+        preventDefault: () => {},
+      } as React.FormEvent<HTMLFormElement>;
+      void handleSubmit(syntheticEvent)
+    }
+  }
+
+  const isSubmitDisabled = !apiKey.trim() || isLoading || localLoading
+
+  return (
+    <Dialog open={open} onOpenChange={onOpenChange}>
+      <DialogContent 
+        className="sm:max-w-md" 
+        showCloseButton={false} // Prevent closing without auth
+      >
+        <DialogHeader>
+          <DialogTitle className="flex items-center gap-2">
+            <Key className="h-5 w-5" />
+            Authentication Required
+          </DialogTitle>
+          <DialogDescription>
+            Please enter your management API key to access the Llamactl dashboard.
+          </DialogDescription>
+        </DialogHeader>
+
+        <form onSubmit={(e) => { void handleSubmit(e) }}>
+          <div className="grid gap-4 py-4">
+            {/* Error Display */}
+            {error && (
+              <div className="flex items-center gap-2 p-3 bg-destructive/10 border border-destructive/20 rounded-lg">
+                <AlertCircle className="h-4 w-4 text-destructive flex-shrink-0" />
+                <span className="text-sm text-destructive">{error}</span>
+              </div>
+            )}
+
+            {/* API Key Input */}
+            <div className="grid gap-2">
+              <Label htmlFor="apiKey">
+                Management API Key <span className="text-red-500">*</span>
+              </Label>
+              <Input
+                id="apiKey"
+                type="password"
+                value={apiKey}
+                onChange={(e) => setApiKey(e.target.value)}
+                onKeyDown={handleKeyDown}
+                placeholder="sk-management-..."
+                disabled={isLoading || localLoading}
+                className={error ? "border-red-500" : ""}
+                autoFocus
+                autoComplete="off"
+              />
+              <p className="text-sm text-muted-foreground">
+                Your management API key is required to access instance management features.
+              </p>
+            </div>
+          </div>
+
+          <DialogFooter className="flex gap-2">
+            <Button
+              type="submit"
+              disabled={isSubmitDisabled}
+              data-testid="login-submit-button"
+            >
+              {(isLoading || localLoading) ? (
+                <>
+                  <Loader2 className="h-4 w-4 animate-spin" />
+                  Authenticating...
+                </>
+              ) : (
+                <>
+                  <Key className="h-4 w-4" />
+                  Login
+                </>
+              )}
+            </Button>
+          </DialogFooter>
+        </form>
+      </DialogContent>
+    </Dialog>
+  )
+}
+
+export default LoginDialog
\ No newline at end of file
diff --git a/webui/src/contexts/AuthContext.tsx b/webui/src/contexts/AuthContext.tsx
new file mode 100644
index 0000000..3d2e7ff
--- /dev/null
+++ b/webui/src/contexts/AuthContext.tsx
@@ -0,0 +1,162 @@
+import { type ReactNode, createContext, useContext, useState, useEffect, useCallback } from 'react'
+
+interface AuthContextState {
+  isAuthenticated: boolean
+  isLoading: boolean
+  apiKey: string | null
+  error: string | null
+}
+
+interface AuthContextActions {
+  login: (apiKey: string) => Promise<void>
+  logout: () => void
+  clearError: () => void
+  validateAuth: () => Promise<boolean>
+}
+
+type AuthContextType = AuthContextState & AuthContextActions
+
+const AuthContext = createContext<AuthContextType | undefined>(undefined)
+
+interface AuthProviderProps {
+  children: ReactNode
+}
+
+const AUTH_STORAGE_KEY = 'llamactl_management_key'
+
+export const AuthProvider = ({ children }: AuthProviderProps) => {
+  const [isAuthenticated, setIsAuthenticated] = useState(false)
+  const [isLoading, setIsLoading] = useState(true)
+  const [apiKey, setApiKey] = useState<string | null>(null)
+  const [error, setError] = useState<string | null>(null)
+
+  // Load auth state from sessionStorage on mount
+  useEffect(() => {
+    const loadStoredAuth = async () => {
+      try {
+        const storedKey = sessionStorage.getItem(AUTH_STORAGE_KEY)
+        if (storedKey) {
+          setApiKey(storedKey)
+          // Validate the stored key
+          const isValid = await validateApiKey(storedKey)
+          if (isValid) {
+            setIsAuthenticated(true)
+          } else {
+            // Invalid key, remove it
+            sessionStorage.removeItem(AUTH_STORAGE_KEY)
+            setApiKey(null)
+          }
+        }
+      } catch (err) {
+        console.error('Error loading stored auth:', err)
+        // Clear potentially corrupted storage
+        sessionStorage.removeItem(AUTH_STORAGE_KEY)
+      } finally {
+        setIsLoading(false)
+      }
+    }
+
+    void loadStoredAuth()
+  }, [])
+
+  // Validate API key by making a test request
+  const validateApiKey = async (key: string): Promise<boolean> => {
+    try {
+      const response = await fetch('/api/v1/instances', {
+        headers: {
+          'Authorization': `Bearer ${key}`,
+          'Content-Type': 'application/json'
+        }
+      })
+      
+      return response.ok
+    } catch (err) {
+      console.error('Auth validation error:', err)
+      return false
+    }
+  }
+
+  const login = useCallback(async (key: string) => {
+    setIsLoading(true)
+    setError(null)
+
+    try {
+      // Validate the provided API key
+      const isValid = await validateApiKey(key)
+      
+      if (!isValid) {
+        throw new Error('Invalid API key')
+      }
+
+      // Store the key and update state
+      sessionStorage.setItem(AUTH_STORAGE_KEY, key)
+      setApiKey(key)
+      setIsAuthenticated(true)
+    } catch (err) {
+      const errorMessage = err instanceof Error ? err.message : 'Authentication failed'
+      setError(errorMessage)
+      throw new Error(errorMessage)
+    } finally {
+      setIsLoading(false)
+    }
+  }, [])
+
+  const logout = useCallback(() => {
+    sessionStorage.removeItem(AUTH_STORAGE_KEY)
+    setApiKey(null)
+    setIsAuthenticated(false)
+    setError(null)
+  }, [])
+
+  const clearError = useCallback(() => {
+    setError(null)
+  }, [])
+
+  const validateAuth = useCallback(async (): Promise<boolean> => {
+    if (!apiKey) return false
+    
+    const isValid = await validateApiKey(apiKey)
+    if (!isValid) {
+      logout()
+    }
+    return isValid
+  }, [apiKey, logout])
+
+  const value: AuthContextType = {
+    isAuthenticated,
+    isLoading,
+    apiKey,
+    error,
+    login,
+    logout,
+    clearError,
+    validateAuth,
+  }
+
+  return (
+    <AuthContext.Provider value={value}>
+      {children}
+    </AuthContext.Provider>
+  )
+}
+
+export const useAuth = (): AuthContextType => {
+  const context = useContext(AuthContext)
+  if (context === undefined) {
+    throw new Error('useAuth must be used within an AuthProvider')
+  }
+  return context
+}
+
+// Helper hook for getting auth headers
+export const useAuthHeaders = (): HeadersInit => {
+  const { apiKey, isAuthenticated } = useAuth()
+  
+  if (!isAuthenticated || !apiKey) {
+    return {}
+  }
+  
+  return {
+    'Authorization': `Bearer ${apiKey}`
+  }
+}
diff --git a/webui/src/contexts/InstancesContext.tsx b/webui/src/contexts/InstancesContext.tsx
index cf37de5..8340bd9 100644
--- a/webui/src/contexts/InstancesContext.tsx
+++ b/webui/src/contexts/InstancesContext.tsx
@@ -1,7 +1,7 @@
-import type { ReactNode } from 'react';
-import { createContext, useContext, useState, useEffect, useCallback } from 'react'
+import { type ReactNode, createContext, useContext, useState, useEffect, useCallback } from 'react'
 import type { CreateInstanceOptions, Instance } from '@/types/instance'
 import { instancesApi } from '@/lib/api'
+import { useAuth } from '@/contexts/AuthContext'
 
 interface InstancesContextState {
   instances: Instance[]
@@ -29,6 +29,7 @@ interface InstancesProviderProps {
 }
 
 export const InstancesProvider = ({ children }: InstancesProviderProps) => {
+  const { isAuthenticated, isLoading: authLoading } = useAuth()
   const [instancesMap, setInstancesMap] = useState<Map<string, Instance>>(new Map())
   const [loading, setLoading] = useState(true)
   const [error, setError] = useState<string | null>(null)
@@ -41,6 +42,11 @@ export const InstancesProvider = ({ children }: InstancesProviderProps) => {
   }, [])
 
   const fetchInstances = useCallback(async () => {
+    if (!isAuthenticated) {
+      setLoading(false)
+      return
+    }
+
     try {
       setLoading(true)
       setError(null)
@@ -57,7 +63,7 @@ export const InstancesProvider = ({ children }: InstancesProviderProps) => {
     } finally {
       setLoading(false)
     }
-  }, [])
+  }, [isAuthenticated])
 
   const updateInstanceInMap = useCallback((name: string, updates: Partial<Instance>) => {
     setInstancesMap(prev => {
@@ -154,9 +160,19 @@ export const InstancesProvider = ({ children }: InstancesProviderProps) => {
     }
   }, [])
 
+  // Only fetch instances when auth is ready and user is authenticated
   useEffect(() => {
-    fetchInstances()
-  }, [fetchInstances])
+    if (!authLoading) {
+      if (isAuthenticated) {
+        void fetchInstances()
+      } else {
+        // Clear instances when not authenticated
+        setInstancesMap(new Map())
+        setLoading(false)
+        setError(null)
+      }
+    }
+  }, [authLoading, isAuthenticated, fetchInstances])
 
   const value: InstancesContextType = {
     instances,
diff --git a/webui/src/lib/api.ts b/webui/src/lib/api.ts
index 66eb63c..25aba18 100644
--- a/webui/src/lib/api.ts
+++ b/webui/src/lib/api.ts
@@ -10,18 +10,31 @@ async function apiCall<T>(
 ): Promise<T> {
   const url = `${API_BASE}${endpoint}`;
 
-  // Prepare headers
-  const headers: HeadersInit = {
+  // Get auth token from sessionStorage (same as AuthContext)
+  const storedKey = sessionStorage.getItem('llamactl_management_key');
+
+  // Prepare headers with auth
+  const headers: Record<string, string> = {
     "Content-Type": "application/json",
-    ...options.headers,
+    ...(options.headers as Record<string, string>),
   };
 
+  // Add auth header if available
+  if (storedKey) {
+    headers['Authorization'] = `Bearer ${storedKey}`;
+  }
+
   try {
     const response = await fetch(url, {
       ...options,
       headers,
     });
 
+    // Handle authentication errors
+    if (response.status === 401) {
+      throw new Error('Authentication required');
+    }
+
     if (!response.ok) {
       // Try to get error message from response
       let errorMessage = `HTTP ${response.status}`;
@@ -47,7 +60,7 @@ async function apiCall<T>(
       const text = await response.text();
       return text as T;
     } else {
-      const data = await response.json();
+      const data = await response.json() as T;
       return data;
     }
   } catch (error) {
diff --git a/webui/src/main.tsx b/webui/src/main.tsx
index ec2cd1d..ab046c2 100644
--- a/webui/src/main.tsx
+++ b/webui/src/main.tsx
@@ -3,11 +3,14 @@ import ReactDOM from 'react-dom/client'
 import App from './App'
 import { InstancesProvider } from './contexts/InstancesContext'
 import './index.css'
+import { AuthProvider } from './contexts/AuthContext'
 
 ReactDOM.createRoot(document.getElementById('root')!).render(
   <React.StrictMode>
-    <InstancesProvider>
-      <App />
-    </InstancesProvider>
+    <AuthProvider>
+      <InstancesProvider>
+        <App />
+      </InstancesProvider>
+    </AuthProvider>
   </React.StrictMode>,
 )
\ No newline at end of file