From e512667e65bb0b23fa16fd62b7ac7659d62d2997 Mon Sep 17 00:00:00 2001
From: Benjamin Blanchard <bblanchard@carnegielearning.com>
Date: Thu, 28 May 2026 15:03:59 -0400
Subject: [PATCH] make return value of updateState consistent with other
 experiment endpoints

---
 .../api/controllers/ExperimentController.ts   |  2 +-
 .../src/api/services/ExperimentService.ts     | 34 +++++++++++--------
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/packages/backend/src/api/controllers/ExperimentController.ts b/packages/backend/src/api/controllers/ExperimentController.ts
index 7586193d3..226b75d79 100644
--- a/packages/backend/src/api/controllers/ExperimentController.ts
+++ b/packages/backend/src/api/controllers/ExperimentController.ts
@@ -1217,7 +1217,7 @@ export class ExperimentController {
     experiment: AssignmentStateUpdateValidator,
     @CurrentUser() currentUser: UserDTO,
     @Req() request: AppRequest
-  ): Promise<any> {
+  ): Promise<ExperimentDTO> {
     return this.experimentService.updateState(
       experiment.experimentId,
       experiment.state,
diff --git a/packages/backend/src/api/services/ExperimentService.ts b/packages/backend/src/api/services/ExperimentService.ts
index 22a962851..97cc32351 100644
--- a/packages/backend/src/api/services/ExperimentService.ts
+++ b/packages/backend/src/api/services/ExperimentService.ts
@@ -476,11 +476,8 @@ export class ExperimentService {
     logger: UpgradeLogger,
     scheduleDate?: Date,
     entityManager?: EntityManager
-  ): Promise<Experiment> {
-    const oldExperiment = await this.experimentRepository.findOne({
-      where: { id: experimentId },
-      relations: ['stateTimeLogs', 'partitions', 'queries', 'queries.metric'],
-    });
+  ): Promise<ExperimentDTO> {
+    const oldExperiment = await this.experimentRepository.findOneExperiment(experimentId);
     await this.clearExperimentCacheDetail(
       oldExperiment.context[0],
       oldExperiment.partitions.map((partition) => {
@@ -529,7 +526,7 @@ export class ExperimentService {
     // updating the experiment and stateTimeLog
     const stateTimeLogRepo = entityManager ? entityManager.getRepository(StateTimeLog) : this.stateTimeLogsRepository;
     logger.info({ message: `stateTimeLogDoc =>`, details: stateTimeLogDoc });
-    const [updatedState, updatedStateTimeLog] = await Promise.all([
+    await Promise.all([
       this.experimentRepository.updateState(experimentId, state, scheduleDate, entityManager),
       stateTimeLogRepo.save(stateTimeLogDoc),
     ]);
@@ -537,15 +534,22 @@ export class ExperimentService {
     // updating experiment schedules
     await this.updateExperimentSchedules(experimentId, logger, entityManager);
 
-    return {
-      ...oldExperiment,
-      state: EXPERIMENT_STATE_DISPLAY_NAME_OVERRIDES[updatedState[0].state] || updatedState[0].state,
-      startOn: updatedState[0].startOn,
-      stateTimeLogs: [
-        ...this.transformStateTimeLogs(oldExperiment.stateTimeLogs),
-        this.transformStateTimeLogs([updatedStateTimeLog]),
-      ].flat(),
-    };
+    oldExperiment.state = EXPERIMENT_STATE_DISPLAY_NAME_OVERRIDES[state] || state;
+    if (scheduleDate) {
+      oldExperiment.startOn = scheduleDate;
+    }
+
+    const responseStateTimeLog = {
+      ...stateTimeLogDoc,
+      experiment: { id: oldExperiment.id },
+    } as StateTimeLog;
+
+    oldExperiment.stateTimeLogs = [
+      ...this.transformStateTimeLogs(oldExperiment.stateTimeLogs),
+      this.transformStateTimeLogs([responseStateTimeLog])[0],
+    ];
+
+    return this.reducedConditionPayload(this.formattingPayload(this.formattingConditionPayload(oldExperiment)));
   }
 
   public async verifyExperiments(