fix(flags): set org/project context before flag fetch (#42123)

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

apps/studio/pages/_app.tsx

@@ -31,7 +31,7 @@ import timezone from 'dayjs/plugin/timezone'
 import utc from 'dayjs/plugin/utc'
 import Head from 'next/head'
 import { NuqsAdapter } from 'nuqs/adapters/next/pages'
-import { ErrorInfo, useCallback } from 'react'
+import { type ComponentProps, ErrorInfo, useCallback } from 'react'
 import { ErrorBoundary } from 'react-error-boundary'
 
 import {
@@ -53,6 +53,7 @@ import { GlobalErrorBoundaryState } from 'components/ui/ErrorBoundary/GlobalErro
 import { useRootQueryClient } from 'data/query-client'
 import { customFont, sourceCodePro } from 'fonts'
 import { useCustomContent } from 'hooks/custom-content/useCustomContent'
+import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization'
 import { AuthProvider } from 'lib/auth'
 import { API_URL, BASE_PATH, IS_PLATFORM, useDefaultProvider } from 'lib/constants'
 import { ProfileProvider } from 'lib/profile'
@@ -67,6 +68,19 @@ dayjs.extend(timezone)
 dayjs.extend(relativeTime)
 dayjs.extend(duration)
 
+const FeatureFlagProviderWithOrgContext = ({
+  children,
+  ...props
+}: ComponentProps<typeof FeatureFlagProvider>) => {
+  const { data: selectedOrganization } = useSelectedOrganizationQuery({ enabled: IS_PLATFORM })
+
+  return (
+    <FeatureFlagProvider {...props} organizationSlug={selectedOrganization?.slug ?? undefined}>
+      {children}
+    </FeatureFlagProvider>
+  )
+}
+
 loader.config({
   // [Joshen] Attempt for offline support/bypass ISP issues is to store the assets required for monaco
   // locally. We're however, only storing the assets which we need (based on what the network tab loads
@@ -124,7 +138,7 @@ function CustomApp({ Component, pageProps }: AppPropsWithLayout) {
         <NuqsAdapter>
           <HydrationBoundary state={pageProps.dehydratedState}>
             <AuthProvider>
-              <FeatureFlagProvider
+              <FeatureFlagProviderWithOrgContext
                 API_URL={API_URL}
                 enabled={IS_PLATFORM}
                 getConfigCatFlags={getConfigCatFlags}
@@ -180,7 +194,7 @@ function CustomApp({ Component, pageProps }: AppPropsWithLayout) {
                     <ReactQueryDevtools initialIsOpen={false} buttonPosition="bottom-left" />
                   )}
                 </ProfileProvider>
-              </FeatureFlagProvider>
+              </FeatureFlagProviderWithOrgContext>
             </AuthProvider>
           </HydrationBoundary>
         </NuqsAdapter>

packages/common/feature-flags.tsx

@@ -1,7 +1,7 @@
 'use client'
 
 import { FlagValues } from 'flags/react'
-import { createContext, PropsWithChildren, useContext, useEffect, useState } from 'react'
+import { createContext, PropsWithChildren, useContext, useEffect, useRef, useState } from 'react'
 
 import { components } from 'api-types'
 import { useAuth } from './auth'
@@ -9,13 +9,27 @@ import { getFlags as getDefaultConfigCatFlags } from './configcat'
 import { hasConsented } from './consent-state'
 import { get, post } from './fetchWrappers'
 import { ensurePlatformSuffix } from './helpers'
+import { useParams } from './hooks'
 
 type TrackFeatureFlagVariables = components['schemas']['TelemetryFeatureFlagBody']
 export type CallFeatureFlagsResponse = components['schemas']['TelemetryCallFeatureFlagsResponse']
 
-export async function getFeatureFlags(API_URL: string) {
+export async function getFeatureFlags(
+  API_URL: string,
+  options: { organizationSlug?: string; projectRef?: string } = {}
+) {
   try {
-    const data = await get(`${ensurePlatformSuffix(API_URL)}/telemetry/feature-flags`)
+    const url = new URL(`${ensurePlatformSuffix(API_URL)}/telemetry/feature-flags`)
+
+    if (options.organizationSlug) {
+      url.searchParams.set('organization_slug', options.organizationSlug)
+    }
+
+    if (options.projectRef) {
+      url.searchParams.set('project_ref', options.projectRef)
+    }
+
+    const data = await get(url.toString())
     return data as CallFeatureFlagsResponse
   } catch (error: any) {
     if (error.message.includes('Failed to fetch')) {
@@ -61,12 +75,16 @@ function getCookies() {
 export const FeatureFlagProvider = ({
   API_URL,
   enabled = true,
+  organizationSlug,
+  projectRef,
   getConfigCatFlags,
   children,
 }: PropsWithChildren<{
   API_URL?: string
   /** Accepts either `boolean` which controls all feature flags or `{ cc: boolean, ph: boolean }` for individual providers */
   enabled?: boolean | { cc: boolean; ph: boolean }
+  organizationSlug?: string
+  projectRef?: string
   /** Custom fetcher for ConfigCat flags if passing in custom attributes */
   getConfigCatFlags?: (
     userEmail?: string
@@ -74,6 +92,10 @@ export const FeatureFlagProvider = ({
 }>) => {
   const { session, isLoading } = useAuth()
   const userEmail = session?.user?.email
+  const params = useParams()
+  const resolvedOrganizationSlug = organizationSlug ?? params.slug
+  const resolvedProjectRef = projectRef ?? params.ref
+  const lastSentGroupContextRef = useRef<string | null>(null)
 
   const [store, setStore] = useState<FeatureFlagContextType>({
     API_URL,
@@ -85,6 +107,35 @@ export const FeatureFlagProvider = ({
   useEffect(() => {
     let mounted = true
 
+    async function ensureGroupContext() {
+      if (!API_URL) return
+
+      const userId = session?.user?.id
+      if (!userId) return
+      if (!hasConsented()) return
+
+      if (!resolvedOrganizationSlug && !resolvedProjectRef) return
+
+      const contextKey = [userId, resolvedOrganizationSlug ?? '', resolvedProjectRef ?? ''].join(
+        '|'
+      )
+      if (lastSentGroupContextRef.current === contextKey) return
+
+      try {
+        await post(
+          `${ensurePlatformSuffix(API_URL)}/telemetry/identify`,
+          {
+            user_id: userId,
+            ...(resolvedOrganizationSlug && { organization_slug: resolvedOrganizationSlug }),
+            ...(resolvedProjectRef && { project_ref: resolvedProjectRef }),
+          },
+          { headers: { Version: '2' } }
+        )
+
+        lastSentGroupContextRef.current = contextKey
+      } catch {}
+    }
+
     async function processFlags() {
       if (!enabled || isLoading) return
 
@@ -96,7 +147,15 @@ export const FeatureFlagProvider = ({
 
       // Run both async operations in parallel
       const [flags, flagValues] = await Promise.all([
-        loadPHFlags ? getFeatureFlags(API_URL) : Promise.resolve({}),
+        loadPHFlags
+          ? (async () => {
+              await ensureGroupContext()
+              return getFeatureFlags(API_URL, {
+                organizationSlug: resolvedOrganizationSlug,
+                projectRef: resolvedProjectRef,
+              })
+            })()
+          : Promise.resolve({}),
         loadCCFlags
           ? typeof getConfigCatFlags === 'function'
             ? getConfigCatFlags(userEmail)
@@ -140,7 +199,16 @@ export const FeatureFlagProvider = ({
     return () => {
       mounted = false
     }
-  }, [enabled, isLoading, userEmail])
+  }, [
+    enabled,
+    isLoading,
+    userEmail,
+    API_URL,
+    session?.user?.id,
+    resolvedOrganizationSlug,
+    resolvedProjectRef,
+    getConfigCatFlags,
+  ])
 
   return (
     <FeatureFlagContext.Provider value={store}>