Merge pull request #932 from PAIR-code/fixes

Correct cohort deletion re-routing + subsequent cohort creation

Author: vivtsai

frontend/src/components/experiment_dashboard/cohort_list.ts

@@ -134,8 +134,8 @@ export class Component extends MobxLitElement {
   }
 
   private renderHeader() {
-    const getCohortName = (length: number) => {
-      return this.experimentManager.getNextCohortName(length);
+    const getCohortName = (offset: number) => {
+      return this.experimentManager.getNextCohortName(offset);
     };
 
     return html`
@@ -155,9 +155,7 @@ export class Component extends MobxLitElement {
               ?loading=${this.experimentManager.isWritingCohort}
               @click=${() => {
                 this.analyticsService.trackButtonClick(ButtonClick.COHORT_ADD);
-                const name = getCohortName(
-                  this.experimentManager.cohortList.length,
-                );
+                const name = getCohortName(0);
                 this.experimentManager.createCohort({}, name);
               }}
             >
@@ -170,12 +168,11 @@ export class Component extends MobxLitElement {
               variant="tonal"
               ?loading=${this.experimentManager.isWritingCohort}
               @click=${() => {
-                const numCohorts = this.experimentManager.cohortList.length;
                 [...Array(5).keys()].forEach((key) => {
                   this.analyticsService.trackButtonClick(
                     ButtonClick.COHORT_ADD,
                   );
-                  const name = getCohortName(numCohorts + key);
+                  const name = getCohortName(key);
                   this.experimentManager.createCohort({}, name);
                 });
               }}
@@ -189,12 +186,11 @@ export class Component extends MobxLitElement {
               variant="tonal"
               ?loading=${this.experimentManager.isWritingCohort}
               @click=${() => {
-                const numCohorts = this.experimentManager.cohortList.length;
                 [...Array(10).keys()].forEach((key) => {
                   this.analyticsService.trackButtonClick(
                     ButtonClick.COHORT_ADD,
                   );
-                  const name = getCohortName(numCohorts + key);
+                  const name = getCohortName(key);
                   this.experimentManager.createCohort({}, name);
                 });
               }}

frontend/src/components/experiment_dashboard/cohort_settings_dialog.ts

@@ -51,11 +51,13 @@ export class CohortSettingsDialog extends MobxLitElement {
           <pr-button
             color="error"
             variant="tonal"
-            @click=${() => {
+            @click=${async () => {
               this.analyticsService.trackButtonClick(ButtonClick.COHORT_DELETE);
-              this.experimentManager.deleteCohort(
+              await this.experimentManager.deleteCohort(
                 this.experimentManager.cohortEditing!.id,
               );
+              this.experimentManager.setCurrentCohortId(undefined);
+              this.experimentManager.setShowCohortList(true, true);
             }}
           >
             Delete cohort

frontend/src/services/experiment.manager.ts

@@ -369,14 +369,27 @@ export class ExperimentManager extends Service {
   }
 
   // Return name for next cohort based on number of cohorts
-  getNextCohortName(numCohorts = this.cohortList.length) {
+  getNextCohortName(offset = 0) {
     const hasTransfer = this.sp.experimentService.stages.find(
       (stage) => stage.kind === StageKind.TRANSFER,
     );
-    if (numCohorts > 0 || !hasTransfer) {
-      return `Cohort ${String(numCohorts).padStart(2, '0')}`;
+
+    if (this.cohortList.length === 0 && hasTransfer) {
+      return 'Lobby';
     }
-    return 'Lobby';
+
+    let maxNum = -1;
+    this.cohortList.forEach((cohort) => {
+      const match = cohort.metadata.name.match(/^Cohort (\d+)$/);
+      if (match) {
+        const num = parseInt(match[1], 10);
+        if (!isNaN(num) && num > maxNum) {
+          maxNum = num;
+        }
+      }
+    });
+
+    return `Cohort ${String(maxNum + 1 + offset).padStart(2, '0')}`;
   }
 
   isFullCohort(cohort: CohortConfig) {
@@ -464,14 +477,13 @@ export class ExperimentManager extends Service {
           'alerts',
         ),
         (snapshot) => {
-          let changedDocs = snapshot.docChanges().map((change) => change.doc);
-          if (changedDocs.length === 0) {
-            changedDocs = snapshot.docs;
-          }
-
-          changedDocs.forEach((doc) => {
-            const data = doc.data() as AlertMessage;
-            this.alertMap[data.id] = data;
+          snapshot.docChanges().forEach((change) => {
+            if (change.type === 'removed') {
+              delete this.alertMap[change.doc.id];
+            } else {
+              const data = change.doc.data() as AlertMessage;
+              this.alertMap[data.id] = data;
+            }
           });
         },
       ),
@@ -487,21 +499,25 @@ export class ExperimentManager extends Service {
           'cohorts',
         ),
         (snapshot) => {
-          let changedDocs = snapshot.docChanges().map((change) => change.doc);
-          if (changedDocs.length === 0) {
-            changedDocs = snapshot.docs;
-          }
-
-          changedDocs.forEach((doc) => {
-            const data = doc.data() as CohortConfig;
-            this.cohortMap[doc.id] = data;
-            this.currentCohortId = doc.id;
+          snapshot.docChanges().forEach((change) => {
+            if (change.type === 'removed') {
+              delete this.cohortMap[change.doc.id];
+              if (this.currentCohortId === change.doc.id) {
+                this.currentCohortId = undefined;
+              }
+            } else {
+              const data = change.doc.data() as CohortConfig;
+              this.cohortMap[change.doc.id] = data;
+              if (!this.currentCohortId) {
+                this.currentCohortId = change.doc.id;
+              }
+            }
           });
 
           // If multiple cohorts, show cohort list
-          if (changedDocs.length > 1) {
+          if (Object.keys(this.cohortMap).length > 1) {
             this.setShowCohortList(true, true);
-          } else if (changedDocs.length === 0) {
+          } else if (Object.keys(this.cohortMap).length === 0) {
             this.setShowCohortEditor(true, true);
           }
 
@@ -523,17 +539,16 @@ export class ExperimentManager extends Service {
           where('currentStatus', '!=', ParticipantStatus.DELETED),
         ),
         (snapshot) => {
-          let changedDocs = snapshot.docChanges().map((change) => change.doc);
-          if (changedDocs.length === 0) {
-            changedDocs = snapshot.docs;
-          }
-
-          changedDocs.forEach((doc) => {
-            const data = {
-              agentConfig: null,
-              ...doc.data(),
-            } as ParticipantProfileExtended;
-            this.participantMap[doc.id] = data;
+          snapshot.docChanges().forEach((change) => {
+            if (change.type === 'removed') {
+              delete this.participantMap[change.doc.id];
+            } else {
+              const data = {
+                agentConfig: null,
+                ...change.doc.data(),
+              } as ParticipantProfileExtended;
+              this.participantMap[change.doc.id] = data;
+            }
           });
 
           this.isParticipantsLoading = false;
@@ -554,17 +569,16 @@ export class ExperimentManager extends Service {
           where('currentStatus', '!=', ParticipantStatus.DELETED),
         ),
         (snapshot) => {
-          let changedDocs = snapshot.docChanges().map((change) => change.doc);
-          if (changedDocs.length === 0) {
-            changedDocs = snapshot.docs;
-          }
-
-          changedDocs.forEach((doc) => {
-            const data = {
-              agentConfig: null,
-              ...doc.data(),
-            } as MediatorProfileExtended;
-            this.mediatorMap[doc.id] = data;
+          snapshot.docChanges().forEach((change) => {
+            if (change.type === 'removed') {
+              delete this.mediatorMap[change.doc.id];
+            } else {
+              const data = {
+                agentConfig: null,
+                ...change.doc.data(),
+              } as MediatorProfileExtended;
+              this.mediatorMap[change.doc.id] = data;
+            }
           });
 
           this.isMediatorsLoading = false;
@@ -586,14 +600,13 @@ export class ExperimentManager extends Service {
           ),
         ),
         (snapshot) => {
-          let changedDocs = snapshot.docChanges().map((change) => change.doc);
-          if (changedDocs.length === 0) {
-            changedDocs = snapshot.docs;
-          }
-
-          changedDocs.forEach((doc) => {
-            const data = doc.data() as AgentPersonaConfig;
-            this.agentPersonaMap[doc.id] = data;
+          snapshot.docChanges().forEach((change) => {
+            if (change.type === 'removed') {
+              delete this.agentPersonaMap[change.doc.id];
+            } else {
+              const data = change.doc.data() as AgentPersonaConfig;
+              this.agentPersonaMap[change.doc.id] = data;
+            }
           });
 
           this.isAgentsLoading = false;
@@ -710,7 +723,7 @@ export class ExperimentManager extends Service {
         cohortId,
       },
     );
-    this.loadExperimentData(this.experimentId);
+
     this.cohortEditing = undefined;
     return response;
   }