Added onboarding rollback on frontend, when onboarding ingestion fails

frontend/app/api/mutations/useOnboardingRollbackMutation.ts

@@ -0,0 +1,44 @@
+import {
+  type UseMutationOptions,
+  useMutation,
+  useQueryClient,
+} from "@tanstack/react-query";
+
+interface OnboardingRollbackResponse {
+  message: string;
+}
+
+export const useOnboardingRollbackMutation = (
+  options?: Omit<
+    UseMutationOptions<OnboardingRollbackResponse, Error, void>,
+    "mutationFn"
+  >,
+) => {
+  const queryClient = useQueryClient();
+
+  async function rollbackOnboarding(): Promise<OnboardingRollbackResponse> {
+    const response = await fetch("/api/onboarding/rollback", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+      },
+    });
+
+    if (!response.ok) {
+      const error = await response.json();
+      throw new Error(error.error || "Failed to rollback onboarding");
+    }
+
+    return response.json();
+  }
+
+  return useMutation({
+    mutationFn: rollbackOnboarding,
+    onSettled: () => {
+      // Invalidate settings query to refetch updated data
+      queryClient.invalidateQueries({ queryKey: ["settings"] });
+    },
+    ...options,
+  });
+};
+

frontend/app/onboarding/_components/onboarding-card.tsx

@@ -3,12 +3,13 @@
 import { useQueryClient } from "@tanstack/react-query";
 import { AnimatePresence, motion } from "framer-motion";
 import { X } from "lucide-react";
-import { useEffect, useState } from "react";
+import { useEffect, useRef, useState } from "react";
 import { toast } from "sonner";
 import {
   type OnboardingVariables,
   useOnboardingMutation,
 } from "@/app/api/mutations/useOnboardingMutation";
+import { useOnboardingRollbackMutation } from "@/app/api/mutations/useOnboardingRollbackMutation";
 import { useGetSettingsQuery } from "@/app/api/queries/useGetSettingsQuery";
 import { useGetTasksQuery } from "@/app/api/queries/useGetTasksQuery";
 import type { ProviderHealthResponse } from "@/app/api/queries/useProviderHealthQuery";
@@ -170,12 +171,32 @@ const OnboardingCard = ({
 
   const [error, setError] = useState<string | null>(null);
 
+  // Track which tasks we've already handled to prevent infinite loops
+  const handledFailedTasksRef = useRef<Set<string>>(new Set());
+
   // Query tasks to track completion
   const { data: tasks } = useGetTasksQuery({
     enabled: currentStep !== null, // Only poll when onboarding has started
     refetchInterval: currentStep !== null ? 1000 : false, // Poll every 1 second during onboarding
   });
 
+  // Rollback mutation
+  const rollbackMutation = useOnboardingRollbackMutation({
+    onSuccess: () => {
+      console.log("Onboarding rolled back successfully");
+      // Reset to provider selection step
+      // Error message is already set before calling mutate
+      setCurrentStep(null);
+    },
+    onError: (error) => {
+      console.error("Failed to rollback onboarding", error);
+      // Preserve existing error message if set, otherwise show rollback error
+      setError((prevError) => prevError || `Failed to rollback: ${error.message}`);
+      // Still reset to provider selection even if rollback fails
+      setCurrentStep(null);
+    },
+  });
+
   // Monitor tasks and call onComplete when all tasks are done
   useEffect(() => {
     if (currentStep === null || !tasks || !isEmbedding) {
@@ -190,11 +211,86 @@ const OnboardingCard = ({
         task.status === "processing",
     );
 
+    // Check if any file failed in completed tasks
+    const completedTasks = tasks.filter(
+      (task) => task.status === "completed"
+    );
+
+    // Check if any completed task has at least one failed file
+    const taskWithFailedFile = completedTasks.find((task) => {
+      // Must have files object
+      if (!task.files || typeof task.files !== "object") {
+        return false;
+      }
+
+      const fileEntries = Object.values(task.files);
+      
+      // Must have at least one file
+      if (fileEntries.length === 0) {
+        return false;
+      }
+
+      // Check if any file has failed status
+      const hasFailedFile = fileEntries.some(
+        (file) => file.status === "failed" || file.status === "error"
+      );
+
+      return hasFailedFile;
+    });
+
+    // If any file failed, show error and jump back one step (like onboardingMutation.onError)
+    // Only handle if we haven't already handled this task
+    if (
+      taskWithFailedFile && 
+      !rollbackMutation.isPending && 
+      !isCompleted &&
+      !handledFailedTasksRef.current.has(taskWithFailedFile.task_id)
+    ) {
+      console.error("File failed in task, jumping back one step", taskWithFailedFile);
+      
+      // Mark this task as handled to prevent infinite loops
+      handledFailedTasksRef.current.add(taskWithFailedFile.task_id);
+      
+      // Extract error messages from failed files
+      const errorMessages: string[] = [];
+      if (taskWithFailedFile.files) {
+        Object.values(taskWithFailedFile.files).forEach((file) => {
+          if ((file.status === "failed" || file.status === "error") && file.error) {
+            errorMessages.push(file.error);
+          }
+        });
+      }
+      
+      // Also check task-level error
+      if (taskWithFailedFile.error) {
+        errorMessages.push(taskWithFailedFile.error);
+      }
+      
+      // Use the first error message, or a generic message if no errors found
+      const errorMessage = errorMessages.length > 0
+        ? errorMessages[0]
+        : "Sample data file failed to ingest. Please try again with a different configuration.";
+      
+      // Set error message and jump back one step (exactly like onboardingMutation.onError)
+      setError(errorMessage);
+      setCurrentStep(totalSteps);
+      // Jump back one step after 1 second (go back to the step before ingestion)
+      // For embedding: totalSteps is 4, ingestion is step 3, so go back to step 2
+      // For LLM: totalSteps is 3, ingestion is step 2, so go back to step 1
+      setTimeout(() => {
+        // Go back to the step before the last step (which is ingestion)
+        const previousStep = totalSteps > 1 ? totalSteps - 2 : 0;
+        setCurrentStep(previousStep);
+      }, 1000);
+      return;
+    }
+
     // If no active tasks and we've started onboarding, complete it
     if (
       (!activeTasks || (activeTasks.processed_files ?? 0) > 0) &&
       tasks.length > 0 &&
-      !isCompleted
+      !isCompleted &&
+      !taskWithFailedFile
     ) {
       // Set to final step to show "Done"
       setCurrentStep(totalSteps);
@@ -203,7 +299,7 @@ const OnboardingCard = ({
         onComplete();
       }, 1000);
     }
-  }, [tasks, currentStep, onComplete, isCompleted, isEmbedding, totalSteps]);
+  }, [tasks, currentStep, onComplete, isCompleted, isEmbedding, totalSteps, rollbackMutation]);
 
   // Mutations
   const onboardingMutation = useOnboardingMutation({

frontend/contexts/chat-context.tsx

@@ -11,6 +11,7 @@ import {
   useState,
 } from "react";
 import { ONBOARDING_STEP_KEY } from "@/lib/constants";
+import { useGetSettingsQuery } from "@/app/api/queries/useGetSettingsQuery";
 
 export type EndpointType = "chat" | "langflow";
 
@@ -115,23 +116,32 @@ export function ChatProvider({ children }: ChatProviderProps) {
     useState<KnowledgeFilter | null>(null);
   const [hasChatError, setChatError] = useState(false);
   
-  // Check if onboarding is complete (onboarding step key should be null)
+  // Get settings to check if onboarding was completed (settings.edited)
+  const { data: settings } = useGetSettingsQuery();
+  
+  // Check if onboarding is complete
+  // Onboarding is complete if:
+  // 1. settings.edited is true (backend confirms onboarding was completed)
+  // 2. AND onboarding step key is null (local onboarding flow is done)
   const [isOnboardingComplete, setIsOnboardingComplete] = useState(() => {
     if (typeof window === "undefined") return false;
-    return localStorage.getItem(ONBOARDING_STEP_KEY) === null;
+    // Default to false if settings not loaded yet
+    return false;
   });
 
-  // Sync onboarding completion state with localStorage
+  // Sync onboarding completion state with settings.edited and localStorage
   useEffect(() => {
     const checkOnboarding = () => {
       if (typeof window !== "undefined") {
-        setIsOnboardingComplete(
-          localStorage.getItem(ONBOARDING_STEP_KEY) === null,
-        );
+        // Onboarding is complete if settings.edited is true AND step key is null
+        const stepKeyExists = localStorage.getItem(ONBOARDING_STEP_KEY) !== null;
+        const isEdited = settings?.edited === true;
+        // Complete if edited is true and step key doesn't exist (onboarding flow finished)
+        setIsOnboardingComplete(isEdited && !stepKeyExists);
       }
     };
 
-    // Check on mount
+    // Check on mount and when settings change
     checkOnboarding();
 
     // Listen for storage events (for cross-tab sync)
@@ -140,7 +150,7 @@ export function ChatProvider({ children }: ChatProviderProps) {
     return () => {
       window.removeEventListener("storage", checkOnboarding);
     };
-  }, []);
+  }, [settings?.edited]);
 
   const setOnboardingComplete = useCallback((complete: boolean) => {
     setIsOnboardingComplete(complete);