fix(media-tab): restore wizard step when changing organizations

Author: derianrddev

apps/pro-web/components/routes/workspace/media-tab/ui/brand-kit/wizard/brand-kit-wizard.tsx

@@ -29,8 +29,8 @@ export function BrandKitWizard({ onFinish, onCancel }: BrandKitWizardProps) {
 		setSelectedVibes,
 		setSelectedFonts,
 		setSelectedPalette,
-		brandKitCurrentStep,
 		saveBrandKitStep,
+		loadBrandKitStep,
 	} = useWorkspaceMedia()
 	const { activeOrganization, initialState } = useWorkspace()
 
@@ -41,15 +41,45 @@ export function BrandKitWizard({ onFinish, onCancel }: BrandKitWizardProps) {
 
 	const [isInitialized, setIsInitialized] = useState(false)
 	const previousBrandKitRef = useRef<string>('')
+	const previousOrgRef = useRef<string | null>(null)
 
 	const { currentStep, Next, Prev, close, goTo, lastStep } = useWizard(
 		steps,
 		true,
 	)
 
+	// Detect organization change and reload step + brand kit
+	useEffect(() => {
+		if (!currentOrg?.id) return
+
+		// If organization changed, reload everything for that org
+		if (previousOrgRef.current !== currentOrg.id) {
+			console.log(`🔄 Organization changed to: ${currentOrg.name}`)
+
+			// Load step from localStorage for this specific org
+			const stepNumber = loadBrandKitStep(currentOrg.id)
+			goTo(stepNumber)
+			console.log(`✅ Loaded step ${stepNumber} for org ${currentOrg.name}`)
+
+			previousOrgRef.current = currentOrg.id
+		}
+
+		// Always mark as initialized after org is set
+		if (!isInitialized) {
+			setIsInitialized(true)
+		}
+	}, [currentOrg?.id, currentOrg?.name, goTo, loadBrandKitStep, isInitialized])
+
 	// Restore Brand Kit data when it changes
 	useEffect(() => {
-		if (!brandKit) return
+		if (!brandKit) {
+			// Clear selections if no brand kit
+			setSelectedVibes([])
+			setSelectedFonts([])
+			setSelectedPalette(null)
+			previousBrandKitRef.current = ''
+			return
+		}
 
 		// Prevent unnecessary updates by comparing serialized brandKit
 		const brandKitStr = JSON.stringify(brandKit)
@@ -60,6 +90,8 @@ export function BrandKitWizard({ onFinish, onCancel }: BrandKitWizardProps) {
 		// Update vibes
 		if (brandKit.vibes?.length > 0) {
 			setSelectedVibes(brandKit.vibes)
+		} else {
+			setSelectedVibes([])
 		}
 
 		// Update fonts
@@ -72,32 +104,21 @@ export function BrandKitWizard({ onFinish, onCancel }: BrandKitWizardProps) {
 				fonts.push({ family: brandKit.fonts.secondary, role: 'secondary' })
 			}
 			setSelectedFonts(fonts)
+		} else {
+			setSelectedFonts([])
 		}
 
 		// Update palette
 		if (brandKit.palette) {
 			setSelectedPalette(brandKit.palette)
+		} else {
+			setSelectedPalette(null)
 		}
 
 		previousBrandKitRef.current = brandKitStr
 		console.log('✅ Brand Kit selections updated')
 	}, [brandKit, setSelectedVibes, setSelectedFonts, setSelectedPalette])
 
-	// Restore step from LocalStorage/State ONCE on mount
-	useEffect(() => {
-		if (isInitialized) return
-
-		// Restore saved step if available
-		if (brandKitCurrentStep && brandKitCurrentStep > 0) {
-			goTo(brandKitCurrentStep)
-			console.log(
-				`✅ Restored to step ${brandKitCurrentStep} from LocalStorage`,
-			)
-		}
-
-		setIsInitialized(true)
-	}, [brandKitCurrentStep, goTo, isInitialized])
-
 	// Save step to LocalStorage on every change
 	useEffect(() => {
 		if (!isInitialized || !currentOrg?.id) return

apps/pro-web/lib/hooks/use-workspace-media.tsx

@@ -12,7 +12,6 @@ import type {
 import type { BrandKitData, ReferenceImagePayload } from '@/types/media.types'
 import type { OrganizationData } from '@/types/thread.types'
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
-import { useLocalStorage } from 'mb-lib'
 import { useSession } from 'next-auth/react'
 import {
 	type ReactNode,
@@ -111,6 +110,7 @@ interface MediaWorkspaceActions {
 		organizationId: string,
 	) => Promise<unknown>
 	saveBrandKitStep: (step: number, organizationId: string) => void
+	loadBrandKitStep: (organizationId: string) => number
 	clearBrandKitStep: (organizationId: string) => void
 }
 
@@ -267,7 +267,6 @@ export function WorkspaceMediaProvider({ children }: { children: ReactNode }) {
 
 		// Extract aspect ratio and name
 		const aspectRatio: AspectRatio = size.ratio
-		const targetFormat = size.name
 
 		// Prompt for resizing/reframing
 		const prompt = `
@@ -407,6 +406,14 @@ export function WorkspaceMediaProvider({ children }: { children: ReactNode }) {
 		[],
 	)
 
+	const loadBrandKitStep = useCallback((organizationId: string) => {
+		const key = `brandkit_step_${organizationId}`
+		const savedStep = localStorage.getItem(key)
+		const step = savedStep ? Number.parseInt(savedStep, 10) : 0
+		setBrandKitCurrentStep(step)
+		return step
+	}, [])
+
 	const clearBrandKitStep = useCallback((organizationId: string) => {
 		setBrandKitCurrentStep(0)
 		const key = `brandkit_step_${organizationId}`
@@ -472,6 +479,7 @@ export function WorkspaceMediaProvider({ children }: { children: ReactNode }) {
 				setSelectedFonts,
 				setSelectedPalette,
 				saveBrandKit,
+				loadBrandKitStep,
 				saveBrandKitStep,
 				clearBrandKitStep,
 			}}