Merge remote-tracking branch 'origin/staging' into fix/docs-generator-hidden-params

Author: waleedlatif1

.agents/skills/add-feature-flag/SKILL.md

@@ -1,16 +1,16 @@
 ---
 name: add-feature-flag
-description: Add a runtime feature flag (AppConfig-backed on prod, secret fallback off-prod), global by default or optionally gated by org id, user id, or platform admin
+description: Add a runtime feature flag (AppConfig-backed on prod, secret fallback off-prod), global by default or optionally gated by workspace id, org id, user id, or platform admin
 argument-hint: <flag-name>
 ---
 
 # Add Feature Flag Skill
 
-You add a **runtime feature flag** to Sim that can change on prod with no redeploy (AWS AppConfig). Prefer a global on/off flag unless the rollout actually needs per-organization, per-user, or platform-admin targeting. When AppConfig isn't the source of truth, the flag falls back to a single **secret** (on/off only).
+You add a **runtime feature flag** to Sim that can change on prod with no redeploy (AWS AppConfig). Prefer a global on/off flag unless the rollout actually needs per-workspace, per-organization, per-user, or platform-admin targeting. When AppConfig isn't the source of truth, the flag falls back to a single **secret** (on/off only).
 
 ## When to use this vs `env-flags.ts`
 
-- **Feature flag** (`@/lib/core/config/feature-flags.ts`): runtime global on/off by default, optionally scoped by `userId`/`orgId`/admin. This skill.
+- **Feature flag** (`@/lib/core/config/feature-flags.ts`): runtime global on/off by default, optionally scoped by `workspaceId`/`userId`/`orgId`/admin. This skill.
 - **Env flag** (`@/lib/core/config/env-flags.ts`): deploy-time capability/environment detection (`isProd`, `isHosted`, `isBillingEnabled`). A module-load boolean. **Do not add gated flags here.**
 
 If the user wants a fixed per-deployment toggle, send them to `env-flags.ts` instead.
@@ -21,10 +21,11 @@ A flag's **gating rule lives only in the hosted AppConfig document**. It is ON f
 
 ```ts
 interface FeatureFlagRule {
-  enabled?: boolean   // global default for everyone
-  orgIds?: string[]   // allowlisted organization ids
-  userIds?: string[]  // allowlisted user ids
-  adminEnabled?: boolean // platform admins (user.role === 'admin')
+  enabled?: boolean       // global default for everyone
+  workspaceIds?: string[] // allowlisted workspace ids
+  orgIds?: string[]       // allowlisted organization ids
+  userIds?: string[]      // allowlisted user ids
+  adminEnabled?: boolean  // platform admins (user.role === 'admin')
 }
 ```
 
@@ -34,10 +35,10 @@ Critically, **none of this is expressible in code** — gating (especially `admi
 
 1. **Confirm the granularity before editing code.** If the user has not already specified it, stop and ask:
 
-   > Should `<flag-name>` be a global on/off flag (recommended), or does it need rollout targeting by organization, user, and/or platform admin?
+   > Should `<flag-name>` be a global on/off flag (recommended), or does it need rollout targeting by workspace, organization, user, and/or platform admin?
 
-   - Recommend **global**. Do not infer scoped gating merely because the call site already has a user or organization id.
-   - If the user chooses scoped gating but does not name the dimensions, ask which of organization, user, and platform admin it needs. Wire only the selected dimensions.
+   - Recommend **global**. Do not infer scoped gating merely because the call site already has a workspace, user, or organization id.
+   - If the user chooses scoped gating but does not name the dimensions, ask which of workspace, organization, user, and platform admin it needs. Wire only the selected dimensions.
    - If the user wants a fixed per-deployment toggle rather than a runtime AppConfig flag, use `env-flags.ts` instead.
 
 2. **Define the flag.** Add one entry to the `FEATURE_FLAGS` registry in `apps/sim/lib/core/config/feature-flags.ts`. Each entry is the flag's whole definition — name (kebab-case key), `description`, and the `fallback` secret consulted when AppConfig isn't the source of truth (truthy ⇒ on globally):
@@ -51,7 +52,7 @@ Critically, **none of this is expressible in code** — gating (especially `admi
    }
    ```
 
-   `fallback` is the env/secret key (typed as `keyof typeof env`), so add `<FLAG_SECRET>` to `apps/sim/lib/core/config/env.ts` first (and the deployment's secret store) — it won't typecheck otherwise. Do **not** add org/user/admin defaults here — that gating exists only in AppConfig. Adding the entry makes `<flag-name>` a valid `FeatureFlagName`.
+   `fallback` is the env/secret key (typed as `keyof typeof env`), so add `<FLAG_SECRET>` to `apps/sim/lib/core/config/env.ts` first (and the deployment's secret store) — it won't typecheck otherwise. Do **not** add workspace/org/user/admin defaults here — that gating exists only in AppConfig. Adding the entry makes `<flag-name>` a valid `FeatureFlagName`.
 
 3. **Gate the call site at the chosen granularity.** For the recommended global mode, pass no context:
 
@@ -70,17 +71,17 @@ Critically, **none of this is expressible in code** — gating (especially `admi
    ```ts
    import { isFeatureEnabled } from '@/lib/core/config/feature-flags'
 
-   if (await isFeatureEnabled('<flag-name>', { userId, orgId })) {
+   if (await isFeatureEnabled('<flag-name>', { workspaceId, userId, orgId })) {
      // gated behavior
    }
    ```
 
-   - Organization targeting uses `orgId`; user and platform-admin targeting require `userId`.
+   - Workspace targeting uses `workspaceId`; organization targeting uses `orgId`; user and platform-admin targeting require `userId`.
    - Missing ids are fine — a clause with no matching id is skipped; with no `userId`, the admin clause resolves to `false` without a DB read.
    - Admin routes that already know the caller is an admin may pass `{ userId, isAdmin: true }` to skip the role lookup.
    - **Client/UI flags:** resolve server-side (in a server component, route, or loader) and pass the boolean down as a prop. There is no client AppConfig.
 
-4. **(Prod) configure in AppConfig.** The infra `feature-flags` profile schema is permissive, so a new flag needs **no infra change**. Operators add the flag to the hosted `feature-flags` document using `enabled` for global rollout or only the selected `orgIds`/`userIds`/`adminEnabled` clauses for scoped rollout, then start a `sim-<env>-fast` deployment (see the AppConfig runbook in the infra README — same flow as `access-control`). The fallback secret only applies when AppConfig is disabled.
+4. **(Prod) configure in AppConfig.** The infra `feature-flags` profile schema is permissive, so a new flag needs **no infra change**. Operators add the flag to the hosted `feature-flags` document using `enabled` for global rollout or only the selected `workspaceIds`/`orgIds`/`userIds`/`adminEnabled` clauses for scoped rollout, then start a `sim-<env>-fast` deployment (see the AppConfig runbook in the infra README — same flow as `access-control`). The fallback secret only applies when AppConfig is disabled.
 
 5. **Test.** Add a case to `apps/sim/lib/core/config/feature-flags.test.ts` that matches the chosen granularity. For a global flag, exercise `isFeatureEnabled('<flag-name>')` with an AppConfig `enabled` rule and toggle the fallback secret for the off-AppConfig path. For scoped rollout, cover only the selected clauses and mock `isPlatformAdmin` when testing `adminEnabled`.
 
@@ -90,6 +91,6 @@ Critically, **none of this is expressible in code** — gating (especially `admi
 
 - Flag keys are `kebab-case`.
 - Never read flags via raw `fetch` or a new AppConfig client — always go through `isFeatureEnabled` / `getFeatureFlags`.
-- Never bake gating into code. The fallback is a single boolean secret; org/user/admin scoping is AppConfig-only.
+- Never bake gating into code. The fallback is a single boolean secret; workspace/org/user/admin scoping is AppConfig-only.
 - Never add or propagate request context unless the user chose scoped rollout.
 - The admin check reads the DB **replica** (`dbReplica`) and is resolved lazily, so an admin-gated flag adds at most one cheap replica read, and only when `adminEnabled` is the deciding clause.

apps/sim/app/api/workspaces/[id]/permissions/route.test.ts

@@ -73,7 +73,7 @@ function queuePersonalWorkspace(
 ) {
   const workspaceRow = { ownerId: OWNER_ID, billedAccountUserId, organizationId: null }
   queueTableRows(schemaMock.workspace, [workspaceRow])
-  /** The in-transaction re-read of the same row, taken `FOR UPDATE`. */
+  /** The in-transaction re-read of the same row, taken `FOR NO KEY UPDATE`. */
   permissionsMockFns.mockGetWorkspaceWithOwner.mockResolvedValue({
     id: WORKSPACE_ID,
     ...workspaceRow,

apps/sim/lib/billing/storage/payer-transfer.test.ts

@@ -769,7 +769,7 @@ describe('changeOrganizationWorkspaceBilledAccountsInTx', () => {
     expect(returning).toHaveBeenCalledWith({ id: 'workspace.id' })
     expect(select).toHaveBeenCalledWith({ id: 'workspace.id' })
     expect(orderBy).toHaveBeenCalledTimes(1)
-    expect(lock).toHaveBeenCalledWith('update')
+    expect(lock).toHaveBeenCalledWith('no key update')
     expect(lock.mock.invocationCallOrder[0]).toBeLessThan(update.mock.invocationCallOrder[0])
     expect(execute).not.toHaveBeenCalled()
   })

apps/sim/lib/billing/storage/payer-transfer.ts

@@ -125,15 +125,17 @@ async function getExactWorkspaceStorageBytes(tx: DbOrTx, workspaceId: string): P
 /**
  * Locks a payer row and returns its current aggregate. A missing source can be
  * historical drift and is represented as `null`; callers must reject a
- * missing destination.
+ * missing destination. `FOR NO KEY UPDATE` avoids upgrading the implicit
+ * foreign-key `FOR KEY SHARE` this transaction may already hold; see the
+ * module header of `lib/billing/storage/tracking.ts`.
  */
 async function lockStoragePayer(tx: DbOrTx, payer: BillingEntity): Promise<number | null> {
   if (payer.type === 'organization') {
     const [row] = await tx
       .select({ storageUsedBytes: organization.storageUsedBytes })
       .from(organization)
       .where(eq(organization.id, payer.id))
-      .for('update')
+      .for('no key update')
       .limit(1)
     return row?.storageUsedBytes ?? null
   }
@@ -142,7 +144,7 @@ async function lockStoragePayer(tx: DbOrTx, payer: BillingEntity): Promise<numbe
     .select({ storageUsedBytes: userStats.storageUsedBytes })
     .from(userStats)
     .where(eq(userStats.userId, payer.id))
-    .for('update')
+    .for('no key update')
     .limit(1)
   return row?.storageUsedBytes ?? null
 }
@@ -257,7 +259,7 @@ async function lockStoragePayers(
       .from(userStats)
       .where(inArray(userStats.userId, userIds))
       .orderBy(asc(userStats.userId))
-      .for('update')
+      .for('no key update')
     for (const row of rows) {
       usageByKey.set(getPayerKey({ type: 'user', id: row.id }), row.storageUsedBytes)
     }
@@ -269,7 +271,7 @@ async function lockStoragePayers(
       .from(organization)
       .where(inArray(organization.id, organizationIds))
       .orderBy(asc(organization.id))
-      .for('update')
+      .for('no key update')
     for (const row of rows) {
       usageByKey.set(getPayerKey({ type: 'organization', id: row.id }), row.storageUsedBytes)
     }
@@ -371,7 +373,7 @@ export async function changeWorkspaceStoragePayersInTx(
     .from(workspace)
     .where(inArray(workspace.id, workspaceIds))
     .orderBy(asc(workspace.id))
-    .for('update')
+    .for('no key update')
 
   const workspaceById = new Map(lockedWorkspaces.map((row) => [row.id, row]))
   for (const workspaceId of workspaceIds) {
@@ -562,7 +564,7 @@ export async function changeOrganizationWorkspaceBilledAccountsInTx(
       )
     )
     .orderBy(asc(workspace.id))
-    .for('update')
+    .for('no key update')
 
   const rows = await tx
     .update(workspace)
@@ -604,7 +606,7 @@ export async function changeWorkspaceStoragePayerInTx(
     })
     .from(workspace)
     .where(eq(workspace.id, params.workspaceId))
-    .for('update')
+    .for('no key update')
     .limit(1)
 
   if (!lockedWorkspace) {

apps/sim/lib/billing/storage/tracking.test.ts

@@ -13,6 +13,7 @@ const {
   mockMaybeNotifyLimit,
   mockOrderedLockRows,
   mockSql,
+  mockTxFor,
   mockTxFrom,
   mockTxLimit,
   mockTxOrderBy,
@@ -32,6 +33,7 @@ const {
   mockMaybeNotifyLimit: vi.fn(),
   mockOrderedLockRows: { queue: [] as unknown[][] },
   mockSql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })),
+  mockTxFor: vi.fn(),
   mockTxFrom: vi.fn(),
   mockTxLimit: vi.fn(),
   mockTxOrderBy: vi.fn(),
@@ -120,6 +122,42 @@ const ORG_CONTEXT: StorageBillingContext = {
   customStorageLimitGB: null,
 }
 
+const USER_CONTEXT: StorageBillingContext = {
+  workspaceId: 'workspace-1',
+  billedAccountUserId: 'workspace-owner',
+  billingEntity: { type: 'user', id: 'workspace-owner' },
+  plan: 'pro',
+  customStorageLimitGB: null,
+}
+
+/**
+ * Both payer kinds. The workspace lock is shared, but the payer lock branches
+ * to a different table per kind, so a lock-mode regression on only one of them
+ * has to fail a test.
+ */
+const PAYER_CASES = [
+  {
+    label: 'organization',
+    context: ORG_CONTEXT,
+    workspaceRow: {
+      billedAccountUserId: 'workspace-owner',
+      organizationId: 'workspace-org' as string | null,
+      storageUsedBytes: 1_000,
+    },
+    payerLockRows: [{ id: 'workspace-org', storageUsedBytes: 1_000 }],
+  },
+  {
+    label: 'user',
+    context: USER_CONTEXT,
+    workspaceRow: {
+      billedAccountUserId: 'workspace-owner',
+      organizationId: null as string | null,
+      storageUsedBytes: 1_000,
+    },
+    payerLockRows: [{ id: 'workspace-owner', storageUsedBytes: 1_000 }],
+  },
+] as const
+
 beforeAll(() => {
   setEnvFlags({ isBillingEnabled: true })
 })
@@ -138,9 +176,10 @@ describe('workspace storage counter mutations', () => {
 
     mockOrderedLockRows.queue = []
     mockTxSelect.mockReturnValue({ from: mockTxFrom })
+    mockTxFor.mockReturnValue({ limit: mockTxLimit })
     mockTxFrom.mockReturnValue({
       where: vi.fn(() => ({
-        for: vi.fn(() => ({ limit: mockTxLimit })),
+        for: mockTxFor,
         limit: mockTxLimit,
         orderBy: mockTxOrderBy,
       })),
@@ -183,6 +222,39 @@ describe('workspace storage counter mutations', () => {
     expect(mockMaybeNotifyLimit).not.toHaveBeenCalled()
   })
 
+  /**
+   * `FOR UPDATE` on these rows deadlocked in production: `workspace`,
+   * `organization`, and `user_stats` are foreign-key parents, so the calling
+   * transaction already holds an implicit `FOR KEY SHARE` on them from the
+   * billable child row it just wrote, and the stronger lock is an upgrade that
+   * two concurrent uploads take on each other. `FOR NO KEY UPDATE` still
+   * conflicts with itself, so the ledgers stay serialized.
+   */
+  it.each(PAYER_CASES)(
+    'locks the workspace and its $label payer as FOR NO KEY UPDATE',
+    async ({ context, workspaceRow }) => {
+      mockWorkspaceRow.current = { ...workspaceRow }
+
+      await incrementStorageUsageForBillingContextInTx(mockTx as unknown as DbOrTx, context, 100)
+
+      expect(mockTxFor.mock.calls).toEqual([['no key update'], ['no key update']])
+    }
+  )
+
+  it.each(PAYER_CASES)(
+    'locks batched workspace and $label payer ledgers as FOR NO KEY UPDATE',
+    async ({ context, workspaceRow, payerLockRows }) => {
+      mockOrderedLockRows.queue = [[{ id: 'workspace-1', ...workspaceRow }], [...payerLockRows]]
+
+      await applyStorageUsageDeltasInTx(mockTx as unknown as DbOrTx, {
+        workspaceDeltas: [{ context, deltaBytes: 100 }],
+        legacyDeltas: [],
+      })
+
+      expect(mockTxOrderedFor.mock.calls).toEqual([['no key update'], ['no key update']])
+    }
+  )
+
   it('serializes quota admission on the locked payer ledger', async () => {
     mockGetStorageLimitForBillingContext.mockReturnValue(1_050)
     mockTxLimit

apps/sim/lib/billing/storage/tracking.ts

@@ -1,5 +1,20 @@
 /**
  * Storage usage tracking for durable workspace and payer ledgers.
+ *
+ * Every row lock here is `FOR NO KEY UPDATE`, never `FOR UPDATE`. The
+ * `workspace`, `organization`, and `user_stats` rows these transactions lock
+ * are foreign-key parents (49 tables reference `workspace` alone), so any
+ * insert or update of a child row — a `workspace_files` row in this very
+ * transaction — implicitly takes `FOR KEY SHARE` on the parent first. A later
+ * `FOR UPDATE` on the same row is then a lock upgrade, and two concurrent
+ * uploads or deletes in one workspace deadlock on it. `FOR NO KEY UPDATE`
+ * does not conflict with `FOR KEY SHARE`, yet still conflicts with itself and
+ * with `FOR UPDATE`, so writers remain serialized against each other and
+ * against payer transfers. It is exactly the lock a plain `UPDATE` of these
+ * non-key counters takes anyway. The only key columns on these tables are
+ * `workspace.id`, `workspace.inbox_provider_id`, `organization.id`,
+ * `user_stats.id`, and `user_stats.user_id`, and no path under these locks
+ * writes any of them or deletes a locked row.
  */
 
 import { organization, userStats, workspace } from '@sim/db/schema'
@@ -124,6 +139,7 @@ async function mutateStorageUsage(
 
 /**
  * Locks and reads the payer ledger after the workspace row has been locked.
+ * `FOR NO KEY UPDATE` for the reason documented at the top of this module.
  */
 async function lockStorageUsageForMutation(
   tx: DbOrTx,
@@ -134,7 +150,7 @@ async function lockStorageUsageForMutation(
       .select({ storageUsedBytes: organization.storageUsedBytes })
       .from(organization)
       .where(eq(organization.id, billingEntity.id))
-      .for('update')
+      .for('no key update')
       .limit(1)
     if (!row) throw new Error(`Storage payer organization:${billingEntity.id} not found`)
     return row.storageUsedBytes
@@ -144,7 +160,7 @@ async function lockStorageUsageForMutation(
     .select({ storageUsedBytes: userStats.storageUsedBytes })
     .from(userStats)
     .where(eq(userStats.userId, billingEntity.id))
-    .for('update')
+    .for('no key update')
     .limit(1)
   if (!row) throw new Error(`Storage payer user:${billingEntity.id} not found`)
   return row.storageUsedBytes
@@ -242,7 +258,7 @@ export async function applyStorageUsageDeltasInTx(
           .from(workspace)
           .where(inArray(workspace.id, workspaceIds))
           .orderBy(asc(workspace.id))
-          .for('update')
+          .for('no key update')
       : []
   const workspaceById = new Map(lockedWorkspaces.map((row) => [row.id, row]))
 
@@ -318,7 +334,7 @@ export async function applyStorageUsageDeltasInTx(
       .from(userStats)
       .where(inArray(userStats.userId, userIds))
       .orderBy(asc(userStats.userId))
-      .for('update')
+      .for('no key update')
     for (const row of rows) {
       payerUsageByKey.set(getPayerKey({ type: 'user', id: row.id }), row.storageUsedBytes)
     }
@@ -329,7 +345,7 @@ export async function applyStorageUsageDeltasInTx(
       .from(organization)
       .where(inArray(organization.id, organizationIds))
       .orderBy(asc(organization.id))
-      .for('update')
+      .for('no key update')
     for (const row of rows) {
       payerUsageByKey.set(getPayerKey({ type: 'organization', id: row.id }), row.storageUsedBytes)
     }
@@ -439,7 +455,7 @@ async function mutateWorkspaceStorageUsage(
     })
     .from(workspace)
     .where(eq(workspace.id, workspaceId))
-    .for('update')
+    .for('no key update')
     .limit(1)
 
   if (!workspacePayer) {

apps/sim/lib/copilot/tools/handlers/materialize-file.ts

@@ -140,7 +140,8 @@ async function executeSave(
 
     try {
       transition = await db.transaction(async (tx) => {
-        await tx.execute(sql`SELECT 1 FROM workspace WHERE id = ${workspaceId} FOR UPDATE`)
+        /** `FOR NO KEY UPDATE`: see the module header of `lib/billing/storage/tracking.ts`. */
+        await tx.execute(sql`SELECT 1 FROM workspace WHERE id = ${workspaceId} FOR NO KEY UPDATE`)
 
         const [updated] = await tx
           .update(workspaceFiles)