diff --git a/apps/src/dance/lab2/views/DanceView.tsx b/apps/src/dance/lab2/views/DanceView.tsx
index 76be447574f17..d02e245c2b8ba 100644
--- a/apps/src/dance/lab2/views/DanceView.tsx
+++ b/apps/src/dance/lab2/views/DanceView.tsx
@@ -133,6 +133,11 @@ const DanceView: React.FunctionComponent<{
     startOver,
   } = useSources<DanceProjectSources>();
 
+  const sourcesRef = useRef(currentSources);
+  useEffect(() => {
+    sourcesRef.current = currentSources;
+  }, [currentSources]);
+
   const mergeSources = useCallback(
     (patch: Partial<DanceProjectSources>, forceSave = false) => {
       const next = {...sourcesRef.current, ...patch};
@@ -468,13 +473,19 @@ const DanceView: React.FunctionComponent<{
       return;
     }
     // In case there is no song set in the current sources, set it to the default.
-    if (!currentSources.selectedSong) {
+    if (!currentSources.selectedSong && !usingMusicProject) {
       const defaultSong = levelProperties.defaultSong;
       const songToUse =
         defaultSong && songData[defaultSong] ? defaultSong : songKeys[0];
       mergeSources({selectedSong: songToUse});
     }
-  }, [songData, currentSources, mergeSources, levelProperties.defaultSong]);
+  }, [
+    songData,
+    currentSources,
+    mergeSources,
+    levelProperties.defaultSong,
+    usingMusicProject,
+  ]);
 
   // Load the selected song whenever it changes in project sources.
   useEffect(() => {
@@ -579,11 +590,6 @@ const DanceView: React.FunctionComponent<{
 
   const settings = useBlocklySettings();
 
-  const sourcesRef = useRef(currentSources);
-  useEffect(() => {
-    sourcesRef.current = currentSources;
-  }, [currentSources]);
-
   if (isShareView) {
     const musicMetadata = loadedMusicProject
       ? musicProjectPlayer.current?.getMetadata()
diff --git a/apps/src/dance/lab2/views/GenerateDance.tsx b/apps/src/dance/lab2/views/GenerateDance.tsx
index 478e85f2cc79a..0ef9ef1151f01 100644
--- a/apps/src/dance/lab2/views/GenerateDance.tsx
+++ b/apps/src/dance/lab2/views/GenerateDance.tsx
@@ -129,7 +129,7 @@ const GenerateDance: React.FunctionComponent<GenerateCodeProps> = ({
     if (aiGenerateState === 'none' && blockCount > 1) {
       setAiGenerateState('edited');
     } else if (
-      ['editing', 'edited'].includes(aiGenerateState) &&
+      ['listened', 'editing', 'edited'].includes(aiGenerateState) &&
       blockCount <= 1
     ) {
       resetProgram();
@@ -341,8 +341,6 @@ const GenerateDance: React.FunctionComponent<GenerateCodeProps> = ({
             size="s"
             onClick={() => {
               startOver();
-              setAiGenerateState('none');
-              resetProgram();
               analyticsReporter.sendEvent(
                 EVENTS.DANCE_PARTY_GENERATE_CODE_BACK_TO_PROMPT_CLICKED,
                 {levelPath: window.location.pathname}
@@ -419,8 +417,6 @@ const GenerateDance: React.FunctionComponent<GenerateCodeProps> = ({
               size="s"
               onClick={() => {
                 startOver();
-                setAiGenerateState('none');
-                resetProgram();
                 analyticsReporter.sendEvent(
                   EVENTS.DANCE_PARTY_GENERATE_CODE_BACK_TO_PROMPT_CLICKED,
                   {levelPath: window.location.pathname}