From 722ba685ebe0525415d210d8c4b5626ec9c5f60e Mon Sep 17 00:00:00 2001
From: Xiaohan-Tian <157918347+Xiaohan-Tian@users.noreply.github.com>
Date: Sat, 4 Jul 2026 23:41:35 -0700
Subject: [PATCH] fix: persist separator results until project switch

---
 src/components/KGOnePanel.test.tsx |  81 ++++++++++-
 src/components/KGOnePanel.tsx      | 219 +++++++++++++++--------------
 2 files changed, 192 insertions(+), 108 deletions(-)

diff --git a/src/components/KGOnePanel.test.tsx b/src/components/KGOnePanel.test.tsx
index 74a86d7..b895434 100644
--- a/src/components/KGOnePanel.test.tsx
+++ b/src/components/KGOnePanel.test.tsx
@@ -12,6 +12,8 @@ const { mockLocalSeparatorDownload } = vi.hoisted(() => ({
 
 let kgoneEnabled = false;
 let selectedRegionIds: string[] = [];
+let projectName = 'Test Project';
+let savedProjectName = 'Test Project';
 let localModelCached: Record<string, boolean> = {};
 let localSeparationResult: Array<{ name: string; blob: Blob }> = [];
 
@@ -21,7 +23,8 @@ const mockExecuteCommand = vi.fn();
 vi.mock('../stores/projectStore', () => ({
   useProjectStore: () => ({
     selectedRegionIds,
-    projectName: 'Test Project',
+    projectName,
+    savedProjectName,
     bpm: 120,
     keySignature: 'C major',
     timeSignature: { numerator: 4, denominator: 4 },
@@ -127,6 +130,8 @@ describe('KGOnePanel local separator mode', () => {
   beforeEach(() => {
     kgoneEnabled = false;
     selectedRegionIds = [];
+    projectName = 'Test Project';
+    savedProjectName = 'Test Project';
     localModelCached = {};
     localSeparationResult = [
       { name: 'Instrumental', blob: new Blob(['instrumental'], { type: 'audio/wav' }) },
@@ -185,6 +190,7 @@ describe('KGOnePanel local separator mode', () => {
     render(<KGOnePanel isVisible={true} />);
 
     expect(await screen.findByText(/Select an audio region on the timeline/)).toBeInTheDocument();
+    expect(screen.getByRole('button', { name: 'Separate Stems' })).toBeDisabled();
   });
 
   it('renders local separation outputs after processing completes', async () => {
@@ -202,6 +208,79 @@ describe('KGOnePanel local separator mode', () => {
     });
   });
 
+  it('keeps stem results visible and disables separation after the selection is cleared', async () => {
+    localModelCached[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium] = true;
+    selectedRegionIds = ['audio-region-1'];
+
+    const { rerender } = render(<KGOnePanel isVisible={true} />);
+
+    fireEvent.click(await screen.findByRole('button', { name: 'Separate Stems' }));
+
+    await waitFor(() => {
+      expect(screen.getByText('Instrumental')).toBeInTheDocument();
+      expect(screen.getByText('Vocals')).toBeInTheDocument();
+      expect(screen.getByRole('button', { name: 'Import All Stems to Timeline' })).toBeInTheDocument();
+    });
+
+    selectedRegionIds = [];
+    rerender(<KGOnePanel isVisible={true} />);
+
+    expect(screen.getByText('Instrumental')).toBeInTheDocument();
+    expect(screen.getByText('Vocals')).toBeInTheDocument();
+    expect(screen.getByRole('button', { name: 'Import All Stems to Timeline' })).toBeInTheDocument();
+    expect(screen.queryByText(/Select an audio region on the timeline/)).not.toBeInTheDocument();
+    expect(screen.getByRole('button', { name: 'Separate Stems' })).toBeDisabled();
+  });
+
+  it('clears stem results when the saved project changes', async () => {
+    localModelCached[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium] = true;
+    selectedRegionIds = ['audio-region-1'];
+
+    const { rerender } = render(<KGOnePanel isVisible={true} />);
+
+    fireEvent.click(await screen.findByRole('button', { name: 'Separate Stems' }));
+
+    await waitFor(() => {
+      expect(screen.getByText('Instrumental')).toBeInTheDocument();
+    });
+    expect(screen.getAllByText('Separation complete.').length).toBeGreaterThan(0);
+
+    savedProjectName = 'Loaded Project';
+    projectName = 'Loaded Project';
+    selectedRegionIds = [];
+    rerender(<KGOnePanel isVisible={true} />);
+
+    await waitFor(() => {
+      expect(screen.queryByText('Instrumental')).not.toBeInTheDocument();
+      expect(screen.queryByText('Vocals')).not.toBeInTheDocument();
+    });
+    expect(screen.queryAllByText('Separation complete.')).toHaveLength(0);
+    expect(screen.queryByRole('button', { name: 'Import All Stems to Timeline' })).not.toBeInTheDocument();
+    expect(screen.getByText(/Select an audio region on the timeline/)).toBeInTheDocument();
+  });
+
+  it('preserves stem results when only the project name changes', async () => {
+    localModelCached[LOCAL_SEPARATOR_MODEL_IDS.mdxMedium] = true;
+    selectedRegionIds = ['audio-region-1'];
+
+    const { rerender } = render(<KGOnePanel isVisible={true} />);
+
+    fireEvent.click(await screen.findByRole('button', { name: 'Separate Stems' }));
+
+    await waitFor(() => {
+      expect(screen.getByText('Instrumental')).toBeInTheDocument();
+    });
+
+    projectName = 'Renamed Project';
+    selectedRegionIds = [];
+    rerender(<KGOnePanel isVisible={true} />);
+
+    expect(screen.getByText('Instrumental')).toBeInTheDocument();
+    expect(screen.getByText('Vocals')).toBeInTheDocument();
+    expect(screen.getByRole('button', { name: 'Import All Stems to Timeline' })).toBeInTheDocument();
+    expect(screen.queryByText(/Select an audio region on the timeline/)).not.toBeInTheDocument();
+  });
+
   it('uses Demucs defaults and renders four local stem players', async () => {
     localModelCached[LOCAL_SEPARATOR_MODEL_IDS.htdemucs4s] = true;
     localSeparationResult = [
diff --git a/src/components/KGOnePanel.tsx b/src/components/KGOnePanel.tsx
index 64235e6..e4a177e 100644
--- a/src/components/KGOnePanel.tsx
+++ b/src/components/KGOnePanel.tsx
@@ -856,7 +856,7 @@ function countRepaintTracks(sourceTrackName: string): number {
 
 const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
   const { t } = useI18n();
-  const { selectedRegionIds, projectName, bpm, timeSignature, maxBars, refreshProjectState } = useProjectStore();
+  const { selectedRegionIds, projectName, savedProjectName, bpm, timeSignature, maxBars, refreshProjectState } = useProjectStore();
   const localOnlyMode = mode === 'local-separator';
   const availableSeparatorModels = localOnlyMode ? LOCAL_SEPARATOR_MODEL_OPTIONS : SERVER_SEPARATOR_MODELS;
   const [model, setModel] = useState<string>(availableSeparatorModels[0].value);
@@ -910,12 +910,14 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
     setStemAudioUrls([]);
     setGenStatus('idle');
     setGenHint('');
+    setLocalProgressPercent(0);
+    setLocalProgressText('');
     setErrorMsg('');
     setIsImporting(false);
     setImportError('');
     originalRegionRef.current = null;
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [projectName]);
+  }, [savedProjectName]);
 
   useEffect(() => {
     setModel(availableSeparatorModels[0].value);
@@ -967,6 +969,7 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
     }
     return null;
   }, [selectedRegionIds]);
+  const hasStemResults = stemAudioUrls.length > 0;
 
   const getConfiguredLocalSeparatorModelUrl = useCallback(() => {
     const configured = ConfigManager.instance().get(currentLocalModelConfig.download.configKey);
@@ -1450,119 +1453,121 @@ const SeparatorTab: React.FC<{ mode: KGOneMode }> = ({ mode }) => {
         </div>
       )}
 
-      {selectedAudioRegion ? (
-        <>
+      <>
+        {selectedAudioRegion && (
           <div className="kgone-region-info">
             <div className="kgone-region-info-label">{t('kgone.shared.selectedRegion')}</div>
             <div className="kgone-region-info-value">{selectedAudioRegion.region.getName()}</div>
             <div className="kgone-region-info-label" style={{ marginTop: 4 }}>{t('kgone.shared.track')}</div>
             <div className="kgone-region-info-value">{selectedAudioRegion.trackName}</div>
           </div>
+        )}
 
-          <div className="kgone-field">
-            <label className="kgone-label">{t('kgone.separator.field.separationModel')}</label>
-            <select className="kgone-select" value={model} onChange={e => setModel(e.target.value)}>
-              {availableSeparatorModels.map(m => (
-                <option key={m.value} value={m.value}>{getModelLabel(m.value, t)}</option>
-              ))}
-            </select>
-          </div>
-
-          {localOnlyMode && (
-            <Expander label={t('kgone.shared.advancedSettings')}>
-              <div className="kgone-field">
-                <label className="kgone-label">{t('kgone.separator.field.chunkDuration')}</label>
-                <input
-                  className="kgone-input"
-                  aria-label={t('kgone.separator.field.chunkDuration')}
-                  type="number"
-                  min={1}
-                  step={1}
-                  value={localChunkDurationSeconds}
-                  onChange={e => setLocalChunkDurationSeconds(e.target.value)}
-                  placeholder={t('kgone.separator.field.chunkDurationPlaceholder')}
-                />
-              </div>
-              <div className="kgone-field">
-                <label className="kgone-label">{t('kgone.separator.field.modelOverlap')}</label>
-                <input
-                  className="kgone-input"
-                  aria-label={t('kgone.separator.field.modelOverlap')}
-                  type="number"
-                  min={0.001}
-                  max={0.999}
-                  step={0.01}
-                  value={localOverlap}
-                  onChange={e => setLocalOverlap(e.target.value)}
-                />
-              </div>
-            </Expander>
-          )}
-
-          {/* Stem audio players — shown once separation is complete */}
-          {stemAudioUrls.length > 0 && (
-            <div className="kgone-stems">
-              {stemAudioUrls.map(stem => (
-                <div key={stem.name} className="kgone-stem-player">
-                  <div className="kgone-label" style={{ marginBottom: 4 }}>{stem.name}</div>
-                  <AudioPlayer
-                    src={stem.url}
-                    dragData={taskIdRef.current ? {
-                      audioFileName: `KGOne_Stem_${stem.name}_${taskIdRef.current}.${localOnlyMode ? 'wav' : 'mp3'}`,
-                    } : undefined}
-                  />
-                </div>
-              ))}
-            </div>
-          )}
-
-          {/* Drag-to-track hint — shown after successful separation */}
-          {genStatus === 'done' && (
-            <div className="kgone-hint" dangerouslySetInnerHTML={{ __html: t('kgone.separator.hint.drag') }} />
-          )}
-
-          {/* Bulk import button — shown after successful separation */}
-          {genStatus === 'done' && stemAudioUrls.length > 0 && (
-            <>
-              <button
-                className="dialog-btn dialog-btn-primary kgone-btn-generate"
-                disabled={isImporting}
-                onClick={handleImportAll}
-                style={{ marginTop: 0 }}
-              >
-                {isImporting && <FaCircleNotch className="kgone-spinner" />}
-                {isImporting ? t('kgone.shared.btn.importing') : t('kgone.separator.btn.importAllStems')}
-              </button>
-              {importError && <div className="kgone-error-msg">{importError}</div>}
-            </>
-          )}
-
-          {/* Separation error message */}
-          {genStatus === 'error' && errorMsg && (
-            <div className="kgone-error-msg">{errorMsg}</div>
-          )}
-
-          <button
-            className="dialog-btn dialog-btn-primary kgone-btn-generate"
-            disabled={isGenerating || (localOnlyMode && !isLocalModelCached)}
-            onClick={handleSeparate}
-          >
-            {isGenerating && <FaCircleNotch className="kgone-spinner" />}
-            {btnLabel()}
-          </button>
-
-          {/* Status hint below button */}
-          {(localOnlyMode ? localProgressText : genHint) && (
-            <div className="kgone-gen-hint">{localOnlyMode ? localProgressText : genHint}</div>
-          )}
-        </>
-      ) : (
-        <div className="kgone-separator-hint">
-          {localOnlyMode && !isLocalModelCached
-            ? t('kgone.separator.hint.noRegion.download', { model: getModelLabel(currentLocalModelConfig.id, t) })
-            : t('kgone.separator.hint.noRegion.select')}
+        <div className="kgone-field">
+          <label className="kgone-label">{t('kgone.separator.field.separationModel')}</label>
+          <select className="kgone-select" value={model} onChange={e => setModel(e.target.value)}>
+            {availableSeparatorModels.map(m => (
+              <option key={m.value} value={m.value}>{getModelLabel(m.value, t)}</option>
+            ))}
+          </select>
         </div>
-      )}
+
+        {localOnlyMode && (
+          <Expander label={t('kgone.shared.advancedSettings')}>
+            <div className="kgone-field">
+              <label className="kgone-label">{t('kgone.separator.field.chunkDuration')}</label>
+              <input
+                className="kgone-input"
+                aria-label={t('kgone.separator.field.chunkDuration')}
+                type="number"
+                min={1}
+                step={1}
+                value={localChunkDurationSeconds}
+                onChange={e => setLocalChunkDurationSeconds(e.target.value)}
+                placeholder={t('kgone.separator.field.chunkDurationPlaceholder')}
+              />
+            </div>
+            <div className="kgone-field">
+              <label className="kgone-label">{t('kgone.separator.field.modelOverlap')}</label>
+              <input
+                className="kgone-input"
+                aria-label={t('kgone.separator.field.modelOverlap')}
+                type="number"
+                min={0.001}
+                max={0.999}
+                step={0.01}
+                value={localOverlap}
+                onChange={e => setLocalOverlap(e.target.value)}
+              />
+            </div>
+          </Expander>
+        )}
+
+        {/* Stem audio players — shown once separation is complete */}
+        {hasStemResults && (
+          <div className="kgone-stems">
+            {stemAudioUrls.map(stem => (
+              <div key={stem.name} className="kgone-stem-player">
+                <div className="kgone-label" style={{ marginBottom: 4 }}>{stem.name}</div>
+                <AudioPlayer
+                  src={stem.url}
+                  dragData={taskIdRef.current ? {
+                    audioFileName: `KGOne_Stem_${stem.name}_${taskIdRef.current}.${localOnlyMode ? 'wav' : 'mp3'}`,
+                  } : undefined}
+                />
+              </div>
+            ))}
+          </div>
+        )}
+
+        {/* Drag-to-track hint — shown after successful separation */}
+        {hasStemResults && (
+          <div className="kgone-hint" dangerouslySetInnerHTML={{ __html: t('kgone.separator.hint.drag') }} />
+        )}
+
+        {/* Bulk import button — shown after successful separation */}
+        {hasStemResults && (
+          <>
+            <button
+              className="dialog-btn dialog-btn-primary kgone-btn-generate"
+              disabled={isImporting}
+              onClick={handleImportAll}
+              style={{ marginTop: 0 }}
+            >
+              {isImporting && <FaCircleNotch className="kgone-spinner" />}
+              {isImporting ? t('kgone.shared.btn.importing') : t('kgone.separator.btn.importAllStems')}
+            </button>
+            {importError && <div className="kgone-error-msg">{importError}</div>}
+          </>
+        )}
+
+        {/* Separation error message */}
+        {genStatus === 'error' && errorMsg && (
+          <div className="kgone-error-msg">{errorMsg}</div>
+        )}
+
+        <button
+          className="dialog-btn dialog-btn-primary kgone-btn-generate"
+          disabled={!selectedAudioRegion || isGenerating || (localOnlyMode && !isLocalModelCached)}
+          onClick={handleSeparate}
+        >
+          {isGenerating && <FaCircleNotch className="kgone-spinner" />}
+          {btnLabel()}
+        </button>
+
+        {/* Status hint below button */}
+        {(localOnlyMode ? localProgressText : genHint) && (
+          <div className="kgone-gen-hint">{localOnlyMode ? localProgressText : genHint}</div>
+        )}
+
+        {!selectedAudioRegion && !hasStemResults && (
+          <div className="kgone-separator-hint">
+            {localOnlyMode && !isLocalModelCached
+              ? t('kgone.separator.hint.noRegion.download', { model: getModelLabel(currentLocalModelConfig.id, t) })
+              : t('kgone.separator.hint.noRegion.select')}
+          </div>
+        )}
+      </>
 
       {!localOnlyMode && (
         <div className="kgone-powered-by">