diff --git a/src/components/exercise-generation/lazy-load-exercise-generation-wrapper-or-null.tsx b/src/components/exercise-generation/lazy-load-exercise-generation-wrapper-or-null.tsx
index 8bc0bc1364..22cea3fcaa 100644
--- a/src/components/exercise-generation/lazy-load-exercise-generation-wrapper-or-null.tsx
+++ b/src/components/exercise-generation/lazy-load-exercise-generation-wrapper-or-null.tsx
@@ -49,16 +49,22 @@ export function LazyLoadExerciseGenerationWrapperOrNull({
       closeWizard={() => {
         setShowWizard(false)
 
-        const url = new URL(window.location.href)
-        // Delete existing query param
-        url.searchParams.delete('showAiWizard')
-        // Update URL without reloading the page
-        router
-          .replace(url.pathname + url.search, undefined, {
-            shallow: true,
-          })
-          .then(() => void null)
-          .catch(() => void null)
+        const referrer = router.query.referrer
+        if (referrer === 'exercise-folder') {
+          // If we come from an exercise folder, we want to transition back!
+          router.back()
+        } else {
+          const url = new URL(window.location.href)
+          // Delete existing query param
+          url.searchParams.delete('showAiWizard')
+          // Update URL without reloading the page
+          router
+            .replace(url.pathname + url.search, undefined, {
+              shallow: true,
+            })
+            .then(() => void null)
+            .catch(() => void null)
+        }
       }}
       subject={subject}
       topic={topic}
diff --git a/src/components/user-tools/foldout-author-menus/author-tools.tsx b/src/components/user-tools/foldout-author-menus/author-tools.tsx
index b6fc6666c9..3a9918d15d 100644
--- a/src/components/user-tools/foldout-author-menus/author-tools.tsx
+++ b/src/components/user-tools/foldout-author-menus/author-tools.tsx
@@ -352,13 +352,13 @@ export function AuthorTools({ tools, entityId, data }: AuthorToolsProps) {
                 title={
                   loggedInStrings.ai.exerciseGeneration.buttonTitleSingular
                 }
-                href={`/entity/create/${UuidType.Exercise}/${data.id}?showAiWizard=`}
+                href={`/entity/create/${UuidType.Exercise}/${data.id}?showAiWizard=&referrer=exercise-folder`}
                 icon={faWandSparkles}
               />
               <SubItem
                 key="ai-group-exercise"
                 title={loggedInStrings.ai.exerciseGeneration.buttonTitle}
-                href={`/entity/create/${entityType}/${data.id}?showAiWizard=`}
+                href={`/entity/create/${entityType}/${data.id}?showAiWizard=&referrer=exercise-folder`}
                 icon={faWandMagicSparkles}
               />
             </Fragment>