fix(webapp): stop saving global flags from unsetting the locked ones

apps/webapp/app/routes/admin.feature-flags.tsx

@@ -5,17 +5,18 @@ import { json } from "@remix-run/server-runtime";
 import { typedjson, useTypedLoaderData } from "remix-typedjson";
 import { z } from "zod";
 import { LockClosedIcon } from "@heroicons/react/20/solid";
-import { boundedIn, prisma } from "~/db.server";
+import { prisma } from "~/db.server";
 import { env } from "~/env.server";
 import { dashboardAction, dashboardLoader } from "~/services/routeBuilders/dashboardBuilder";
 import {
   FEATURE_FLAG,
   GLOBAL_LOCKED_FLAGS,
+  type FeatureFlagKey,
   type FlagControlType,
   getAllFlagControlTypes,
   validatePartialFeatureFlags,
 } from "~/v3/featureFlags";
-import { flags as getGlobalFlags } from "~/v3/featureFlags.server";
+import { flags as getGlobalFlags, replaceGlobalFeatureFlags } from "~/v3/featureFlags.server";
 import { featuresForRequest } from "~/features.server";
 import { Button } from "~/components/primitives/Buttons";
 import { Callout } from "~/components/primitives/Callout";
@@ -87,7 +88,13 @@ export const action = dashboardAction(
       return json({ error: "Invalid JSON body" }, { status: 400 });
     }
 
-    const payloadSchema = z.object({ flags: z.record(z.unknown()) });
+    const payloadSchema = z.object({
+      flags: z.record(z.unknown()),
+      // The page only submits the flags it is managing, so an omitted key is ambiguous for the
+      // locked flags: this says whether the admin unlocked them and is therefore authoritative
+      // over them too.
+      unlockLockedFlags: z.boolean().optional(),
+    });
     const parsed = payloadSchema.safeParse(body);
     if (!parsed.success) {
       return json({ error: "Invalid payload" }, { status: 400 });
@@ -116,39 +123,12 @@ export const action = dashboardAction(
       );
     }
 
-    const validatedFlags = validationResult.data as Record<string, unknown>;
-    const controlTypes = getAllFlagControlTypes();
-    const catalogKeys = Object.keys(controlTypes);
-
-    const keysToDelete: string[] = [];
-    const upsertOps: ReturnType<typeof prisma.featureFlag.upsert>[] = [];
-
-    for (const key of catalogKeys) {
-      if (key in validatedFlags) {
-        upsertOps.push(
-          prisma.featureFlag.upsert({
-            where: { key },
-            create: { key, value: validatedFlags[key] as any },
-            update: { value: validatedFlags[key] as any },
-          })
-        );
-      } else {
-        // On cloud, never delete locked flags (they're not in the payload
-        // because the UI doesn't include them). Locally, delete everything
-        // the user didn't include - full control.
-        const isProtected = isManagedCloud && GLOBAL_LOCKED_FLAGS.includes(key);
-        if (!isProtected) {
-          keysToDelete.push(key);
-        }
-      }
-    }
-
-    await prisma.$transaction([
-      ...upsertOps,
-      ...(keysToDelete.length > 0
-        ? [prisma.featureFlag.deleteMany({ where: { key: { in: boundedIn(keysToDelete) } } })]
-        : []),
-    ]);
+    await replaceGlobalFeatureFlags(prisma, {
+      requestedFlags: validationResult.data as Record<string, unknown>,
+      catalogKeys: Object.keys(getAllFlagControlTypes()) as FeatureFlagKey[],
+      isManagedCloud,
+      unlockLockedFlags: parsed.data.unlockLockedFlags ?? false,
+    });
 
     return json({ success: true });
   }
@@ -213,7 +193,7 @@ export default function AdminFeatureFlagsRoute() {
   };
 
   const handleSave = () => {
-    saveFetcher.submit(JSON.stringify({ flags: values }), {
+    saveFetcher.submit(JSON.stringify({ flags: values, unlockLockedFlags: unlocked }), {
       method: "POST",
       encType: "application/json",
     });

apps/webapp/app/v3/featureFlags.server.ts

@@ -1,11 +1,12 @@
 import { type z } from "zod";
 import type { PrismaClient } from "@trigger.dev/database";
-import { prisma, type PrismaClientOrTransaction } from "~/db.server";
+import { $transaction, boundedIn, prisma, type PrismaClientOrTransaction } from "~/db.server";
 import {
   FEATURE_FLAG,
   type FeatureFlagCatalogSchema,
   type FeatureFlagKey,
   FeatureFlagCatalog,
+  GLOBAL_LOCKED_FLAGS,
 } from "~/v3/featureFlags";
 import { stampMintKindFlip } from "~/v3/runOpsMigration/mintFlipGrace";
 
@@ -220,3 +221,55 @@ export async function applyGlobalMintKindFlip(
     return makeSetMultipleFlags(tx)(stamped);
   });
 }
+
+/**
+ * Replace-semantics write for the global admin flags page: catalog keys present in
+ * `requestedFlags` are upserted, catalog keys absent from it are deleted.
+ *
+ * A locked flag absent from the payload means the page never offered it for editing, not that
+ * the admin unset it, so it survives the sweep. Only a self-hosted page that says it unlocked
+ * them can delete one.
+ */
+export async function replaceGlobalFeatureFlags(
+  client: PrismaClient,
+  params: {
+    requestedFlags: Record<string, unknown>;
+    catalogKeys: FeatureFlagKey[];
+    isManagedCloud: boolean;
+    unlockLockedFlags: boolean;
+  }
+): Promise<void> {
+  const canDeleteLocked = params.unlockLockedFlags && !params.isManagedCloud;
+  const toUpsert: { key: FeatureFlagKey; value: unknown }[] = [];
+  const keysToDelete: string[] = [];
+
+  for (const key of params.catalogKeys) {
+    if (key in params.requestedFlags) {
+      toUpsert.push({ key, value: params.requestedFlags[key] });
+    } else if (canDeleteLocked || !GLOBAL_LOCKED_FLAGS.includes(key)) {
+      keysToDelete.push(key);
+    }
+  }
+
+  const applied = await $transaction(client, "replaceGlobalFeatureFlags", async (tx) => {
+    for (const { key, value } of toUpsert) {
+      await tx.featureFlag.upsert({
+        where: { key },
+        create: { key, value: value as any },
+        update: { value: value as any },
+      });
+    }
+
+    if (keysToDelete.length > 0) {
+      await tx.featureFlag.deleteMany({ where: { key: { in: boundedIn(keysToDelete) } } });
+    }
+
+    return true;
+  });
+
+  // The helper resolves undefined instead of throwing when Prisma errors are swallowed. This
+  // write deletes flags, so treat a transaction that did not run as a failure the caller sees.
+  if (!applied) {
+    throw new Error("replaceGlobalFeatureFlags: transaction did not complete");
+  }
+}

apps/webapp/test/globalFeatureFlagsLockedFlags.test.ts

@@ -0,0 +1,112 @@
+// With "Unlock read-only flags" off, the page strips GLOBAL_LOCKED_FLAGS from its payload, so an
+// omitted locked key means "the UI never offered it", not "the admin unset it".
+import type { PrismaClient } from "@trigger.dev/database";
+import { postgresTest } from "@internal/testcontainers";
+import { describe, expect, vi } from "vitest";
+import { FEATURE_FLAG, FeatureFlagCatalog, type FeatureFlagKey } from "~/v3/featureFlags";
+import { makeSetMultipleFlags, replaceGlobalFeatureFlags } from "~/v3/featureFlags.server";
+
+vi.setConfig({ testTimeout: 60_000 });
+
+const CATALOG_KEYS = Object.keys(FeatureFlagCatalog) as FeatureFlagKey[];
+const WORKER_GROUP_ID = "clwg000000000000000000000";
+
+async function readFlag(prisma: PrismaClient, key: FeatureFlagKey): Promise<unknown> {
+  const row = await prisma.featureFlag.findFirst({ where: { key }, select: { value: true } });
+  return row?.value;
+}
+
+describe("replaceGlobalFeatureFlags — locked flags the UI never submitted", () => {
+  postgresTest(
+    "keeps defaultWorkerInstanceGroupId when a locked flag is absent from the payload",
+    async ({ prisma }) => {
+      await makeSetMultipleFlags(prisma)({
+        [FEATURE_FLAG.defaultWorkerInstanceGroupId]: WORKER_GROUP_ID,
+        [FEATURE_FLAG.mollifierEnabled]: true,
+      });
+
+      // What the page posts when an admin unsets mollifierEnabled on a self-hosted instance.
+      await replaceGlobalFeatureFlags(prisma, {
+        requestedFlags: {},
+        catalogKeys: CATALOG_KEYS,
+        isManagedCloud: false,
+        unlockLockedFlags: false,
+      });
+
+      expect(await readFlag(prisma, FEATURE_FLAG.defaultWorkerInstanceGroupId)).toBe(
+        WORKER_GROUP_ID
+      );
+      expect(await readFlag(prisma, FEATURE_FLAG.mollifierEnabled)).toBeUndefined();
+    }
+  );
+
+  postgresTest("an unlocked self-hosted page can still unset a locked flag", async ({ prisma }) => {
+    await makeSetMultipleFlags(prisma)({
+      [FEATURE_FLAG.defaultWorkerInstanceGroupId]: WORKER_GROUP_ID,
+    });
+
+    await replaceGlobalFeatureFlags(prisma, {
+      requestedFlags: {},
+      catalogKeys: CATALOG_KEYS,
+      isManagedCloud: false,
+      unlockLockedFlags: true,
+    });
+
+    expect(await readFlag(prisma, FEATURE_FLAG.defaultWorkerInstanceGroupId)).toBeUndefined();
+  });
+
+  postgresTest(
+    "managed cloud keeps locked flags even when unlocking is claimed",
+    async ({ prisma }) => {
+      await makeSetMultipleFlags(prisma)({
+        [FEATURE_FLAG.defaultWorkerInstanceGroupId]: WORKER_GROUP_ID,
+      });
+
+      await replaceGlobalFeatureFlags(prisma, {
+        requestedFlags: {},
+        catalogKeys: CATALOG_KEYS,
+        isManagedCloud: true,
+        unlockLockedFlags: true,
+      });
+
+      expect(await readFlag(prisma, FEATURE_FLAG.defaultWorkerInstanceGroupId)).toBe(
+        WORKER_GROUP_ID
+      );
+    }
+  );
+
+  postgresTest("managed cloud still sweeps ordinary flags it was not sent", async ({ prisma }) => {
+    await makeSetMultipleFlags(prisma)({
+      [FEATURE_FLAG.defaultWorkerInstanceGroupId]: WORKER_GROUP_ID,
+      [FEATURE_FLAG.hasAiAccess]: true,
+    });
+
+    await replaceGlobalFeatureFlags(prisma, {
+      requestedFlags: { [FEATURE_FLAG.mollifierEnabled]: true },
+      catalogKeys: CATALOG_KEYS,
+      isManagedCloud: true,
+      unlockLockedFlags: false,
+    });
+
+    expect(await readFlag(prisma, FEATURE_FLAG.defaultWorkerInstanceGroupId)).toBe(WORKER_GROUP_ID);
+    expect(await readFlag(prisma, FEATURE_FLAG.hasAiAccess)).toBeUndefined();
+    expect(await readFlag(prisma, FEATURE_FLAG.mollifierEnabled)).toBe(true);
+  });
+
+  postgresTest("submitted flags are upserted and omitted ones swept", async ({ prisma }) => {
+    await makeSetMultipleFlags(prisma)({
+      [FEATURE_FLAG.mollifierEnabled]: true,
+      [FEATURE_FLAG.hasAiAccess]: true,
+    });
+
+    await replaceGlobalFeatureFlags(prisma, {
+      requestedFlags: { [FEATURE_FLAG.mollifierEnabled]: false },
+      catalogKeys: CATALOG_KEYS,
+      isManagedCloud: false,
+      unlockLockedFlags: false,
+    });
+
+    expect(await readFlag(prisma, FEATURE_FLAG.mollifierEnabled)).toBe(false);
+    expect(await readFlag(prisma, FEATURE_FLAG.hasAiAccess)).toBeUndefined();
+  });
+});