From 3783ecef801c970c511f0712b40cdf31dc19d6c8 Mon Sep 17 00:00:00 2001
From: Tanishq Gandhi <tangandhi19@gmail.com>
Date: Thu, 8 Oct 2026 12:14:17 -0700
Subject: [PATCH] fix(frontend): clear the picked version when switching
 datasets in the dataset picker

---
 .../dataset-selection-modal.component.spec.ts | 45 +++++++++++++++++++
 .../dataset-selection-modal.component.ts      |  7 ++-
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/frontend/src/app/workspace/component/dataset-selection-modal/dataset-selection-modal.component.spec.ts b/frontend/src/app/workspace/component/dataset-selection-modal/dataset-selection-modal.component.spec.ts
index f42ae00a818..67da1511b84 100644
--- a/frontend/src/app/workspace/component/dataset-selection-modal/dataset-selection-modal.component.spec.ts
+++ b/frontend/src/app/workspace/component/dataset-selection-modal/dataset-selection-modal.component.spec.ts
@@ -295,6 +295,51 @@ describe("DatasetSelectionModalComponent", () => {
       expect(confirmButton().disabled).toBe(false);
     });
 
+    it("disables Select again when another dataset is picked", () => {
+      datasetService.retrieveAccessibleDatasets.mockReturnValue(of([dataset, sharedDataset]));
+      build();
+      selects()[0].triggerEventHandler("ngModelChange", dataset);
+      selects()[1].triggerEventHandler("ngModelChange", version);
+      fixture.detectChanges();
+      expect(confirmButton().disabled).toBe(false);
+
+      datasetService.retrieveDatasetVersionList.mockReturnValue(of([sharedVersion]));
+      selects()[0].triggerEventHandler("ngModelChange", sharedDataset);
+      fixture.detectChanges();
+
+      expect(component.selectedVersion).toBeUndefined();
+      expect(confirmButton().disabled).toBe(true);
+    });
+
+    it("drops the previous file when switching to a dataset with no versions in file mode", () => {
+      datasetService.retrieveAccessibleDatasets.mockReturnValue(of([dataset, sharedDataset]));
+      modalData.fileMode = true;
+      modalData.selectedPath = `/dataset/${OWNER}/myds/v1`;
+      build();
+      fixture.debugElement
+        .query(By.directive(UserDatasetVersionFiletreeComponent))
+        .triggerEventHandler("selectedTreeNode", fileNode);
+      fixture.detectChanges();
+      expect(confirmButton().disabled).toBe(false);
+
+      datasetService.retrieveDatasetVersionList.mockReturnValue(of([]));
+      selects()[0].triggerEventHandler("ngModelChange", sharedDataset);
+      fixture.detectChanges();
+
+      expect(component.selectedPath).toBeUndefined();
+      expect(confirmButton().disabled).toBe(true);
+    });
+
+    it("reselects the saved version when reopened in non-file mode", () => {
+      modalData.selectedPath = `/dataset/${OWNER}/myds/v1`;
+      build();
+
+      expect(component.selectedVersion).toBe(version);
+      expect(confirmButton().disabled).toBe(false);
+      confirmButton().click();
+      expect(modalRef.close).toHaveBeenCalledWith(`/dataset/${OWNER}/myds/v1`);
+    });
+
     it("closes the modal with the composed path when Select is clicked", () => {
       build();
       selects()[0].triggerEventHandler("ngModelChange", dataset);
diff --git a/frontend/src/app/workspace/component/dataset-selection-modal/dataset-selection-modal.component.ts b/frontend/src/app/workspace/component/dataset-selection-modal/dataset-selection-modal.component.ts
index b5768aa685b..8d5d517a8ad 100644
--- a/frontend/src/app/workspace/component/dataset-selection-modal/dataset-selection-modal.component.ts
+++ b/frontend/src/app/workspace/component/dataset-selection-modal/dataset-selection-modal.component.ts
@@ -100,14 +100,17 @@ export class DatasetSelectionModalComponent implements OnInit {
 
   onDatasetChange(versionName?: string) {
     this.fileTree = [];
+    this.selectedVersion = undefined;
+    this.selectedPath = undefined;
     if (this.selectedDataset?.dataset.did !== undefined) {
       this.datasetService
         .retrieveDatasetVersionList(this.selectedDataset.dataset.did)
         .pipe(untilDestroyed(this))
         .subscribe(versions => {
           this.datasetVersions = versions;
-          if (this.data.fileMode) {
-            this.selectedVersion = versions.find(version => version.name === versionName) ?? versions[0];
+          const savedVersion = versions.find(version => version.name === versionName);
+          this.selectedVersion = this.data.fileMode ? savedVersion ?? versions[0] : savedVersion;
+          if (this.selectedVersion) {
             this.onVersionChange();
           }
         });