From 02ad0641c2ce8e339d36515004255166389884a6 Mon Sep 17 00:00:00 2001
From: Daniel Bachhuber <daniel.b@posthog.com>
Date: Fri, 20 Dec 2024 10:52:05 -0800
Subject: [PATCH] fix(experiments): Avoid overriding feature flag payloads
 (#27102)

---
 .../ExperimentView/DistributionTable.tsx      | 13 +++-------
 .../scenes/experiments/experimentLogic.tsx    | 25 +++++++++++++++++++
 .../scenes/feature-flags/featureFlagLogic.ts  |  2 +-
 3 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/frontend/src/scenes/experiments/ExperimentView/DistributionTable.tsx b/frontend/src/scenes/experiments/ExperimentView/DistributionTable.tsx
index caee718efb726..d34901f9716a1 100644
--- a/frontend/src/scenes/experiments/ExperimentView/DistributionTable.tsx
+++ b/frontend/src/scenes/experiments/ExperimentView/DistributionTable.tsx
@@ -23,12 +23,11 @@ import { VariantScreenshot } from './VariantScreenshot'
 
 export function DistributionModal({ experimentId }: { experimentId: Experiment['id'] }): JSX.Element {
     const { experiment, experimentLoading, isDistributionModalOpen } = useValues(experimentLogic({ experimentId }))
-    const { closeDistributionModal, updateExperiment } = useActions(experimentLogic({ experimentId }))
+    const { closeDistributionModal, updateDistributionModal } = useActions(experimentLogic({ experimentId }))
 
     const _featureFlagLogic = featureFlagLogic({ id: experiment.feature_flag?.id ?? null } as FeatureFlagLogicProps)
     const { featureFlag, areVariantRolloutsValid, variantRolloutSum } = useValues(_featureFlagLogic)
-    const { setFeatureFlagFilters, distributeVariantsEqually, saveSidebarExperimentFeatureFlag } =
-        useActions(_featureFlagLogic)
+    const { setFeatureFlagFilters, distributeVariantsEqually } = useActions(_featureFlagLogic)
 
     const handleRolloutPercentageChange = (index: number, value: number | undefined): void => {
         if (!featureFlag?.filters?.multivariate || !value) {
@@ -61,13 +60,7 @@ export function DistributionModal({ experimentId }: { experimentId: Experiment['
                     </LemonButton>
                     <LemonButton
                         onClick={() => {
-                            saveSidebarExperimentFeatureFlag(featureFlag)
-                            updateExperiment({
-                                holdout_id: experiment.holdout_id,
-                                parameters: {
-                                    feature_flag_variants: featureFlag?.filters?.multivariate?.variants ?? [],
-                                },
-                            })
+                            updateDistributionModal(featureFlag)
                             closeDistributionModal()
                         }}
                         type="primary"
diff --git a/frontend/src/scenes/experiments/experimentLogic.tsx b/frontend/src/scenes/experiments/experimentLogic.tsx
index 42f6b2cd37652..1941845241e4c 100644
--- a/frontend/src/scenes/experiments/experimentLogic.tsx
+++ b/frontend/src/scenes/experiments/experimentLogic.tsx
@@ -10,7 +10,12 @@ import { lemonToast } from 'lib/lemon-ui/LemonToast/LemonToast'
 import { featureFlagLogic } from 'lib/logic/featureFlagLogic'
 import { hasFormErrors, toParams } from 'lib/utils'
 import { eventUsageLogic } from 'lib/utils/eventUsageLogic'
+import {
+    indexToVariantKeyFeatureFlagPayloads,
+    variantKeyToIndexFeatureFlagPayloads,
+} from 'scenes/feature-flags/featureFlagLogic'
 import { validateFeatureFlagKey } from 'scenes/feature-flags/featureFlagLogic'
+import { featureFlagsLogic } from 'scenes/feature-flags/featureFlagsLogic'
 import { funnelDataLogic } from 'scenes/funnels/funnelDataLogic'
 import { insightDataLogic } from 'scenes/insights/insightDataLogic'
 import { cleanFilters, getDefaultEvent } from 'scenes/insights/utils/cleanFilters'
@@ -165,6 +170,8 @@ export const experimentLogic = kea<experimentLogicType>([
             ],
             teamLogic,
             ['addProductIntent'],
+            featureFlagsLogic,
+            ['updateFlag'],
         ],
     })),
     actions({
@@ -262,6 +269,7 @@ export const experimentLogic = kea<experimentLogicType>([
         openPrimaryMetricModal: (index: number) => ({ index }),
         closePrimaryMetricModal: true,
         setPrimaryMetricsResultErrors: (errors: any[]) => ({ errors }),
+        updateDistributionModal: (featureFlag: FeatureFlagType) => ({ featureFlag }),
     }),
     reducers({
         experiment: [
@@ -795,6 +803,23 @@ export const experimentLogic = kea<experimentLogicType>([
                 lemonToast.error('Failed to update experiment variant images')
             }
         },
+        updateDistributionModal: async ({ featureFlag }) => {
+            const { created_at, id, ...flag } = featureFlag
+
+            const preparedFlag = indexToVariantKeyFeatureFlagPayloads(flag)
+
+            const savedFlag = await api.update(
+                `api/projects/${values.currentProjectId}/feature_flags/${id}`,
+                preparedFlag
+            )
+
+            const updatedFlag = variantKeyToIndexFeatureFlagPayloads(savedFlag)
+            actions.updateFlag(updatedFlag)
+
+            actions.updateExperiment({
+                holdout_id: values.experiment.holdout_id,
+            })
+        },
     })),
     loaders(({ actions, props, values }) => ({
         experiment: {
diff --git a/frontend/src/scenes/feature-flags/featureFlagLogic.ts b/frontend/src/scenes/feature-flags/featureFlagLogic.ts
index 4a9fd05113d3e..41eac1ddc740b 100644
--- a/frontend/src/scenes/feature-flags/featureFlagLogic.ts
+++ b/frontend/src/scenes/feature-flags/featureFlagLogic.ts
@@ -154,7 +154,7 @@ export const variantKeyToIndexFeatureFlagPayloads = (flag: FeatureFlagType): Fea
     }
 }
 
-const indexToVariantKeyFeatureFlagPayloads = (flag: Partial<FeatureFlagType>): Partial<FeatureFlagType> => {
+export const indexToVariantKeyFeatureFlagPayloads = (flag: Partial<FeatureFlagType>): Partial<FeatureFlagType> => {
     if (flag.filters?.multivariate) {
         const newPayloads: Record<string, JsonType> = {}
         flag.filters.multivariate.variants.forEach(({ key }, index) => {