made assistant message be saved on local storage to save progress on onboarding

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

@@ -1,6 +1,6 @@
 "use client";
 
-import { useEffect, useState } from "react";
+import { useEffect, useRef, useState } from "react";
 import { StickToBottom } from "use-stick-to-bottom";
 import { AssistantMessage } from "@/app/chat/_components/assistant-message";
 import Nudges from "@/app/chat/_components/nudges";
@@ -8,26 +8,79 @@ import { UserMessage } from "@/app/chat/_components/user-message";
 import type { Message } from "@/app/chat/_types/types";
 import OnboardingCard from "@/app/onboarding/_components/onboarding-card";
 import { useChatStreaming } from "@/hooks/useChatStreaming";
+import {
+  ONBOARDING_ASSISTANT_MESSAGE_KEY,
+  ONBOARDING_SELECTED_NUDGE_KEY,
+} from "@/lib/constants";
 
 import { OnboardingStep } from "./onboarding-step";
 import OnboardingUpload from "./onboarding-upload";
 
 export function OnboardingContent({
   handleStepComplete,
+  handleStepBack,
   currentStep,
 }: {
   handleStepComplete: () => void;
+  handleStepBack: () => void;
   currentStep: number;
 }) {
+  const parseFailedRef = useRef(false);
   const [responseId, setResponseId] = useState<string | null>(null);
-  const [selectedNudge, setSelectedNudge] = useState<string>("");
+  const [selectedNudge, setSelectedNudge] = useState<string>(() => {
+    // Retrieve selected nudge from localStorage on mount
+    if (typeof window === "undefined") return "";
+    return localStorage.getItem(ONBOARDING_SELECTED_NUDGE_KEY) || "";
+  });
   const [assistantMessage, setAssistantMessage] = useState<Message | null>(
-    null,
+    () => {
+      // Retrieve assistant message from localStorage on mount
+      if (typeof window === "undefined") return null;
+      const savedMessage = localStorage.getItem(ONBOARDING_ASSISTANT_MESSAGE_KEY);
+      if (savedMessage) {
+        try {
+          const parsed = JSON.parse(savedMessage);
+          // Convert timestamp string back to Date object
+          return {
+            ...parsed,
+            timestamp: new Date(parsed.timestamp),
+          };
+        } catch (error) {
+          console.error("Failed to parse saved assistant message:", error);
+          parseFailedRef.current = true;
+          // Clear corrupted data - will go back a step in useEffect
+          if (typeof window !== "undefined") {
+            localStorage.removeItem(ONBOARDING_ASSISTANT_MESSAGE_KEY);
+            localStorage.removeItem(ONBOARDING_SELECTED_NUDGE_KEY);
+          }
+          return null;
+        }
+      }
+      return null;
+    },
   );
 
+  // Handle parse errors by going back a step
+  useEffect(() => {
+    if (parseFailedRef.current && currentStep >= 2) {
+      handleStepBack();
+    }
+  }, [handleStepBack, currentStep]);
+
   const { streamingMessage, isLoading, sendMessage } = useChatStreaming({
     onComplete: (message, newResponseId) => {
       setAssistantMessage(message);
+      // Save assistant message to localStorage when complete
+      if (typeof window !== "undefined") {
+        try {
+          localStorage.setItem(
+            ONBOARDING_ASSISTANT_MESSAGE_KEY,
+            JSON.stringify(message),
+          );
+        } catch (error) {
+          console.error("Failed to save assistant message to localStorage:", error);
+        }
+      }
       if (newResponseId) {
         setResponseId(newResponseId);
       }
@@ -47,7 +100,15 @@ export function OnboardingContent({
 
   const handleNudgeClick = async (nudge: string) => {
     setSelectedNudge(nudge);
+    // Save selected nudge to localStorage
+    if (typeof window !== "undefined") {
+      localStorage.setItem(ONBOARDING_SELECTED_NUDGE_KEY, nudge);
+    }
     setAssistantMessage(null);
+    // Clear saved assistant message when starting a new conversation
+    if (typeof window !== "undefined") {
+      localStorage.removeItem(ONBOARDING_ASSISTANT_MESSAGE_KEY);
+    }
     setTimeout(async () => {
       await sendMessage({
         prompt: nudge,

frontend/components/chat-renderer.tsx

@@ -18,6 +18,8 @@ import { useChat } from "@/contexts/chat-context";
 import {
   ANIMATION_DURATION,
   HEADER_HEIGHT,
+  ONBOARDING_ASSISTANT_MESSAGE_KEY,
+  ONBOARDING_SELECTED_NUDGE_KEY,
   ONBOARDING_STEP_KEY,
   SIDEBAR_WIDTH,
   TOTAL_ONBOARDING_STEPS,
@@ -81,15 +83,25 @@ export function ChatRenderer({
       // Onboarding is complete - remove from local storage and show layout
       if (typeof window !== "undefined") {
         localStorage.removeItem(ONBOARDING_STEP_KEY);
+        localStorage.removeItem(ONBOARDING_ASSISTANT_MESSAGE_KEY);
+        localStorage.removeItem(ONBOARDING_SELECTED_NUDGE_KEY);
       }
       setShowLayout(true);
     }
   };
 
+  const handleStepBack = () => {
+    if (currentStep > 0) {
+      setCurrentStep(currentStep - 1);
+    }
+  };
+
   const handleSkipOnboarding = () => {
     // Skip onboarding by marking it as complete
     if (typeof window !== "undefined") {
       localStorage.removeItem(ONBOARDING_STEP_KEY);
+      localStorage.removeItem(ONBOARDING_ASSISTANT_MESSAGE_KEY);
+      localStorage.removeItem(ONBOARDING_SELECTED_NUDGE_KEY);
     }
     setShowLayout(true);
   };
@@ -193,6 +205,7 @@ export function ChatRenderer({
               {!showLayout && (
                 <OnboardingContent
                   handleStepComplete={handleStepComplete}
+                  handleStepBack={handleStepBack}
                   currentStep={currentStep}
                 />
               )}

frontend/lib/constants.ts

@@ -34,6 +34,8 @@ export const TOTAL_ONBOARDING_STEPS = 5;
  * Local Storage Keys
  */
 export const ONBOARDING_STEP_KEY = "onboarding_current_step";
+export const ONBOARDING_ASSISTANT_MESSAGE_KEY = "onboarding_assistant_message";
+export const ONBOARDING_SELECTED_NUDGE_KEY = "onboarding_selected_nudge";
 
 export const FILES_REGEX =
   /(?<=I'm uploading a document called ['"])[^'"]+\.[^.]+(?=['"]\. Here is its content:)/;