refactor opt in form and modal

apps/studio/components/interfaces/Organization/GeneralSettings/AIOptInLevelSelector.tsx

@@ -0,0 +1,100 @@
+import { ReactNode } from 'react'
+import { Control } from 'react-hook-form'
+import { cn, FormField_Shadcn_, RadioGroupCard, RadioGroupCardItem } from 'ui'
+import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
+import OptInToOpenAIToggle from './OptInToOpenAIToggle'
+import { AIOptInFormValues } from 'hooks/forms/useAIOptInForm'
+
+interface AIOptInLevelSelectorProps {
+  control: Control<AIOptInFormValues>
+  disabled?: boolean
+  label?: ReactNode
+  layout?: 'horizontal' | 'vertical' | 'flex-row-reverse'
+}
+
+export const AIOptInLevelSelector = ({
+  control,
+  disabled,
+  label,
+  layout = 'vertical',
+}: AIOptInLevelSelectorProps) => {
+  const content = (
+    <div className="space-y-4">
+      {/* Wrap RadioGroupCard with FormField_Shadcn_ providing name and render */}
+      <FormField_Shadcn_
+        control={control}
+        name="aiOptInLevel" // Ensure this name matches the form schema
+        render={({ field }) => (
+          <RadioGroupCard
+            value={field.value}
+            onValueChange={field.onChange}
+            disabled={disabled}
+            className="grid md:grid-cols-3 gap-2"
+          >
+            <RadioGroupCardItem
+              className={cn(layout !== 'flex-row-reverse' && 'w-full')}
+              key="disabled"
+              value="disabled"
+              label={
+                <span className="flex flex-col gap-2">
+                  <span className="text-foreground">Disabled</span>
+                  <span className="text-foreground-light text-xs">
+                    No data is sent to OpenAI, responses will be generic.
+                  </span>
+                </span>
+              }
+            />
+            <RadioGroupCardItem
+              className={cn(layout !== 'flex-row-reverse' && 'w-full')}
+              key="schema"
+              value="schema"
+              label={
+                <span className="flex flex-col gap-2">
+                  <span className="text-foreground">Schema Only</span>
+                  <span className="text-foreground-light text-xs">
+                    Send only your database schema to OpenAI for better responses.
+                  </span>
+                </span>
+              }
+            />
+            <RadioGroupCardItem
+              className={cn(layout !== 'flex-row-reverse' && 'w-full')}
+              key="schema_and_data"
+              value="schema_and_data"
+              label={
+                <span className="flex flex-col gap-2">
+                  <span className="text-foreground">Schema & Sample Data</span>
+                  <span className="text-foreground-light text-xs">
+                    Send schema and SQL query data for the best AI responses.
+                  </span>
+                </span>
+              }
+            />
+          </RadioGroupCard>
+        )}
+      />
+    </div>
+  )
+
+  // Conditional rendering with FormItemLayout remains the same
+  return (
+    <FormItemLayout
+      label={label}
+      description={
+        <>
+          <p className="mb-4">
+            By opting into sending anonymous data, Supabase AI can improve the answers it shows you.
+            This is an organization-wide setting. Select the level of data you are comfortable
+            sharing.
+          </p>
+          <OptInToOpenAIToggle />
+        </>
+      }
+      layout={layout}
+    >
+      {content}
+    </FormItemLayout>
+  )
+
+  return content
+}

apps/studio/components/interfaces/Organization/GeneralSettings/DataPrivacyForm.tsx

@@ -1,185 +1,39 @@
 import { useEffect } from 'react'
-import { useForm } from 'react-hook-form'
-import { zodResolver } from '@hookform/resolvers/zod'
 import * as z from 'zod'
-import { toast } from 'sonner'
-import { useQueryClient } from '@tanstack/react-query'
-import { useParams } from 'common'
-import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization'
-import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
-import { PermissionAction } from '@supabase/shared-types/out/constants'
-import type { ResponseError } from 'types'
-import { useOrganizationUpdateMutation } from 'data/organizations/organization-update-mutation'
-import { invalidateOrganizationsQuery } from 'data/organizations/organizations-query'
-import { OPT_IN_TAGS } from 'lib/constants'
-import {
-  Form_Shadcn_,
-  Card,
-  CardContent,
-  CardFooter,
-  FormField_Shadcn_,
-  RadioGroupCard,
-  RadioGroupCardItem,
-} from 'ui'
+import { Form_Shadcn_, Card, CardContent, CardFooter } from 'ui'
 import { FormActions } from 'components/ui/Forms/FormActions'
-import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
-import OptInToOpenAIToggle from './OptInToOpenAIToggle'
-
-const DataPrivacySchema = z.object({
-  aiOptInLevel: z.enum(['disabled', 'schema', 'schema_and_data'], {
-    required_error: 'AI Opt-in level selection is required',
-  }),
-})
-
-const getAiOptInLevel = (tags: string[] | undefined): 'disabled' | 'schema' | 'schema_and_data' => {
-  const hasSql = tags?.includes(OPT_IN_TAGS.AI_SQL) ?? false
-  const hasData = tags?.includes(OPT_IN_TAGS.AI_DATA ?? 'AI_DATA')
-
-  if (hasSql && hasData) {
-    return 'schema_and_data'
-  } else if (hasSql) {
-    return 'schema'
-  } else {
-    return 'disabled'
-  }
-}
+import { useAIOptInForm } from 'hooks/forms/useAIOptInForm'
+import { AIOptInLevelSelector } from './AIOptInLevelSelector'
 
 const DataPrivacyForm = () => {
-  const selectedOrganization = useSelectedOrganization()
-  const canUpdateOrganization = useCheckPermissions(PermissionAction.UPDATE, 'organizations')
-  const { slug } = useParams()
-  const queryClient = useQueryClient()
-
-  const { mutate: updateOrganization, isLoading: isUpdatingPrivacy } =
-    useOrganizationUpdateMutation()
-
-  const dataPrivacyForm = useForm<z.infer<typeof DataPrivacySchema>>({
-    resolver: zodResolver(DataPrivacySchema),
-    defaultValues: {
-      aiOptInLevel: getAiOptInLevel(selectedOrganization?.opt_in_tags),
-    },
-  })
+  const { form, onSubmit, isUpdating, canUpdateOrganization, currentOptInLevel } = useAIOptInForm()
 
   useEffect(() => {
-    if (selectedOrganization && !isUpdatingPrivacy) {
-      dataPrivacyForm.reset({
-        aiOptInLevel: getAiOptInLevel(selectedOrganization.opt_in_tags),
-      })
-    }
-  }, [selectedOrganization, dataPrivacyForm.reset, isUpdatingPrivacy])
-
-  const onUpdateDataPrivacy = async (values: z.infer<typeof DataPrivacySchema>) => {
-    if (!canUpdateOrganization) {
-      return toast.error('You do not have the required permissions to update this organization')
-    }
-    if (!slug) return console.error('Slug is required')
-
-    const existingOptInTags = selectedOrganization?.opt_in_tags ?? []
-    let updatedOptInTags = existingOptInTags.filter(
-      (tag: string) => tag !== OPT_IN_TAGS.AI_SQL && tag !== (OPT_IN_TAGS.AI_DATA ?? 'AI_DATA')
-    )
-    if (values.aiOptInLevel === 'schema' || values.aiOptInLevel === 'schema_and_data') {
-      updatedOptInTags.push(OPT_IN_TAGS.AI_SQL)
-    }
-    if (values.aiOptInLevel === 'schema_and_data') {
-      updatedOptInTags.push(OPT_IN_TAGS.AI_DATA ?? 'AI_DATA')
-    }
-    updatedOptInTags = [...new Set(updatedOptInTags)]
-
-    updateOrganization(
-      { slug, opt_in_tags: updatedOptInTags },
-      {
-        onSuccess: () => {
-          invalidateOrganizationsQuery(queryClient)
-          toast.success('Successfully updated data privacy settings')
-        },
-        onError: (error: ResponseError) => {
-          toast.error(`Failed to update data privacy settings: ${error.message}`)
-        },
-      }
-    )
-  }
+    form.reset({ aiOptInLevel: currentOptInLevel })
+  }, [currentOptInLevel, form.reset])
 
   const permissionsHelperText = !canUpdateOrganization
     ? "You need additional permissions to manage this organization's settings"
     : undefined
 
   return (
-    <Form_Shadcn_ {...dataPrivacyForm}>
-      <form id="org-privacy-form" onSubmit={dataPrivacyForm.handleSubmit(onUpdateDataPrivacy)}>
+    <Form_Shadcn_ {...form}>
+      <form id="org-privacy-form" onSubmit={form.handleSubmit(onSubmit)}>
         <Card>
           <CardContent className="pt-6">
-            <FormField_Shadcn_
-              control={dataPrivacyForm.control}
-              name="aiOptInLevel"
-              render={({ field }) => (
-                <FormItemLayout
-                  label="Supabase AI Opt-in Level"
-                  layout="flex-row-reverse"
-                  description={
-                    <>
-                      <p className="mb-4">
-                        By opting into sending anonymous data, Supabase AI can improve the answers
-                        it shows you. This is an organization-wide setting. Select the level of data
-                        you are comfortable sharing.
-                      </p>
-                      <OptInToOpenAIToggle />
-                    </>
-                  }
-                >
-                  <RadioGroupCard
-                    value={field.value}
-                    onValueChange={field.onChange}
-                    disabled={!canUpdateOrganization || isUpdatingPrivacy}
-                    className="grid gap-4 sm:grid-cols-3"
-                  >
-                    <RadioGroupCardItem
-                      key="disabled"
-                      value="disabled"
-                      label={
-                        <span className="flex flex-col gap-2">
-                          <span className="text-foreground">Disabled</span>
-                          <span className="text-foreground-light text-xs">
-                            No data is sent to OpenAI, responses will be generic.
-                          </span>
-                        </span>
-                      }
-                    />
-                    <RadioGroupCardItem
-                      key="schema"
-                      value="schema"
-                      label={
-                        <span className="flex flex-col gap-2">
-                          <span className="text-foreground">Schema Only</span>
-                          <span className="text-foreground-light text-xs">
-                            Send only your database schema to OpenAI for better responses.
-                          </span>
-                        </span>
-                      }
-                    />
-                    <RadioGroupCardItem
-                      key="schema_and_data"
-                      value="schema_and_data"
-                      label={
-                        <span className="flex flex-col gap-2">
-                          <span className="text-foreground">Schema & Sample Data</span>
-                          <span className="text-foreground-light text-xs">
-                            Send schema and SQL query data for the best AI responses.
-                          </span>
-                        </span>
-                      }
-                    />
-                  </RadioGroupCard>
-                </FormItemLayout>
-              )}
+            <AIOptInLevelSelector
+              control={form.control}
+              disabled={!canUpdateOrganization || isUpdating}
+              layout="flex-row-reverse"
+              label="Supabase AI Opt-in Level"
             />
           </CardContent>
           <CardFooter className="flex justify-end p-4 md:px-8">
             <FormActions
               form="org-privacy-form"
-              isSubmitting={isUpdatingPrivacy}
-              hasChanges={dataPrivacyForm.formState.isDirty}
-              handleReset={() => dataPrivacyForm.reset()}
+              isSubmitting={isUpdating}
+              hasChanges={form.formState.isDirty}
+              handleReset={() => form.reset()}
               helper={permissionsHelperText}
               disabled={!canUpdateOrganization}
             />

apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx

@@ -20,9 +20,7 @@ import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-que
 import { useTablesQuery } from 'data/tables/tables-query'
 import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
 import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
-import {
-  useOrgAiOptInLevel,
-} from 'hooks/misc/useOrgOptedIntoAi'
+import { useOrgAiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
 import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization'
 import { useSelectedProject } from 'hooks/misc/useSelectedProject'
 import { useFlag } from 'hooks/ui/useFlag'
@@ -40,7 +38,6 @@ import {
   TooltipTrigger,
 } from 'ui'
 import { Admonition, AssistantChatForm, GenericSkeletonLoader } from 'ui-patterns'
-import ConfirmationModal from 'ui-patterns/Dialogs/ConfirmationModal'
 import { ButtonTooltip } from '../ButtonTooltip'
 import DotGrid from '../DotGrid'
 import { AIAssistantChatSelector } from './AIAssistantChatSelector'
@@ -48,6 +45,7 @@ import AIOnboarding from './AIOnboarding'
 import CollapsibleCodeBlock from './CollapsibleCodeBlock'
 import { Message } from './Message'
 import { useAutoScroll } from './hooks'
+import { AIOptInModal } from './AIOptInModal'
 
 const MemoizedMessage = memo(
   ({
@@ -181,7 +179,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   })
 
   const canUpdateOrganization = useCheckPermissions(PermissionAction.UPDATE, 'organizations')
-  const { mutate: updateOrganization, isLoading: isUpdating } = useOrganizationUpdateMutation()
 
   const updateMessage = useCallback(
     ({
@@ -242,30 +239,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     }
   }
 
-  const confirmOptInToShareSchemaData = async () => {
-    if (!canUpdateOrganization) {
-      return toast.error('You do not have the required permissions to update this organization')
-    }
-
-    if (!selectedOrganization?.slug) return console.error('Organization slug is required')
-
-    const existingOptInTags = selectedOrganization?.opt_in_tags ?? []
-
-    const updatedOptInTags = existingOptInTags.includes(OPT_IN_TAGS.AI_SQL)
-      ? existingOptInTags
-      : [...existingOptInTags, OPT_IN_TAGS.AI_SQL]
-
-    updateOrganization(
-      { slug: selectedOrganization?.slug, opt_in_tags: updatedOptInTags },
-      {
-        onSuccess: () => {
-          toast.success('Successfully opted-in')
-          setIsConfirmOptInModalOpen(false)
-        },
-      }
-    )
-  }
-
   const handleClearMessages = () => {
     snap.clearMessages()
     setMessages([])
@@ -631,22 +604,10 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
         </div>
       </div>
 
-      <ConfirmationModal
+      <AIOptInModal
         visible={isConfirmOptInModalOpen}
-        size="large"
-        title="Confirm sending anonymous data to OpenAI"
-        confirmLabel="Confirm"
         onCancel={() => setIsConfirmOptInModalOpen(false)}
-        onConfirm={confirmOptInToShareSchemaData}
-        loading={isUpdating}
-      >
-        <p className="text-sm text-foreground-light mb-4">
-          By opting into sending anonymous data, Supabase AI can improve the answers it shows you.
-          This is an organization-wide setting, and affects all projects in your organization.
-        </p>
-
-        <OptInToOpenAIToggle />
-      </ConfirmationModal>
+      />
     </>
   )
 }

apps/studio/components/ui/AIAssistantPanel/AIOptInModal.tsx

@@ -0,0 +1,72 @@
+import { useEffect } from 'react'
+import {
+  Button,
+  Dialog,
+  DialogContent,
+  DialogFooter,
+  DialogHeader,
+  DialogSection,
+  DialogSectionSeparator,
+  DialogTitle,
+  Form_Shadcn_,
+} from 'ui'
+import { useAIOptInForm } from 'hooks/forms/useAIOptInForm'
+import { AIOptInLevelSelector } from 'components/interfaces/Organization/GeneralSettings/AIOptInLevelSelector'
+
+interface AIOptInModalProps {
+  visible: boolean
+  onCancel: () => void
+}
+
+export const AIOptInModal = ({ visible, onCancel }: AIOptInModalProps) => {
+  const { form, onSubmit, isUpdating, canUpdateOrganization, currentOptInLevel } =
+    useAIOptInForm(onCancel)
+
+  useEffect(() => {
+    if (visible) {
+      form.reset({ aiOptInLevel: currentOptInLevel })
+    }
+  }, [visible, currentOptInLevel, form.reset])
+
+  const onOpenChange = (open: boolean) => {
+    if (!open) {
+      onCancel()
+    }
+  }
+
+  return (
+    <Dialog open={visible} onOpenChange={onOpenChange}>
+      <DialogContent size="xlarge">
+        <Form_Shadcn_ {...form}>
+          <form id="ai-opt-in-form" onSubmit={form.handleSubmit(onSubmit)}>
+            <DialogHeader padding="small">
+              <DialogTitle>Update Supabase Assistant Opt-in Level</DialogTitle>
+            </DialogHeader>
+            <DialogSectionSeparator />
+            <DialogSection className="space-y-4" padding="small">
+              <AIOptInLevelSelector
+                control={form.control}
+                disabled={!canUpdateOrganization || isUpdating}
+              />
+            </DialogSection>
+            <DialogSectionSeparator />
+            <DialogFooter padding="small">
+              <Button type="default" disabled={isUpdating} onClick={onCancel}>
+                Cancel
+              </Button>
+              <Button
+                type="primary"
+                htmlType="submit"
+                form="ai-opt-in-form"
+                loading={isUpdating}
+                disabled={isUpdating || !canUpdateOrganization || !form.formState.isDirty}
+              >
+                Confirm
+              </Button>
+            </DialogFooter>
+          </form>
+        </Form_Shadcn_>
+      </DialogContent>
+    </Dialog>
+  )
+}

apps/studio/hooks/forms/useAIOptInForm.ts

@@ -0,0 +1,86 @@
+import { useForm } from 'react-hook-form'
+import { zodResolver } from '@hookform/resolvers/zod'
+import * as z from 'zod'
+import { toast } from 'sonner'
+import { useQueryClient } from '@tanstack/react-query'
+
+import { PermissionAction } from '@supabase/shared-types/out/constants'
+import type { ResponseError } from 'types'
+import { useOrganizationUpdateMutation } from 'data/organizations/organization-update-mutation'
+import { invalidateOrganizationsQuery } from 'data/organizations/organizations-query'
+import { useSelectedOrganization } from 'hooks/misc/useSelectedOrganization'
+import { useCheckPermissions } from 'hooks/misc/useCheckPermissions'
+import { getAiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
+import { OPT_IN_TAGS } from 'lib/constants'
+
+// Shared schema definition
+export const AIOptInSchema = z.object({
+  aiOptInLevel: z.enum(['disabled', 'schema', 'schema_and_data'], {
+    required_error: 'AI Opt-in level selection is required',
+  }),
+})
+
+export type AIOptInFormValues = z.infer<typeof AIOptInSchema>
+
+/**
+ * Hook to manage the AI Opt-In form state and submission logic.
+ * Optionally takes an onSuccess callback (e.g., to close a modal).
+ */
+export const useAIOptInForm = (onSuccessCallback?: () => void) => {
+  const queryClient = useQueryClient()
+  const selectedOrganization = useSelectedOrganization()
+  const canUpdateOrganization = useCheckPermissions(PermissionAction.UPDATE, 'organizations')
+
+  const { mutate: updateOrganization, isLoading: isUpdating } = useOrganizationUpdateMutation()
+
+  const form = useForm<AIOptInFormValues>({
+    resolver: zodResolver(AIOptInSchema),
+    defaultValues: {
+      aiOptInLevel: getAiOptInLevel(selectedOrganization?.opt_in_tags),
+    },
+  })
+
+  const onSubmit = async (values: AIOptInFormValues) => {
+    if (!canUpdateOrganization) {
+      return toast.error('You do not have the required permissions to update this organization')
+    }
+    if (!selectedOrganization?.slug) {
+      console.error('Organization slug is required')
+      return toast.error('Failed to update settings: Organization not found.')
+    }
+
+    const existingOptInTags = selectedOrganization?.opt_in_tags ?? []
+    let updatedOptInTags = existingOptInTags.filter(
+      (tag: string) => tag !== OPT_IN_TAGS.AI_SQL && tag !== (OPT_IN_TAGS.AI_DATA ?? 'AI_DATA')
+    )
+    if (values.aiOptInLevel === 'schema' || values.aiOptInLevel === 'schema_and_data') {
+      updatedOptInTags.push(OPT_IN_TAGS.AI_SQL)
+    }
+    if (values.aiOptInLevel === 'schema_and_data') {
+      updatedOptInTags.push(OPT_IN_TAGS.AI_DATA ?? 'AI_DATA')
+    }
+    updatedOptInTags = [...new Set(updatedOptInTags)]
+
+    updateOrganization(
+      { slug: selectedOrganization.slug, opt_in_tags: updatedOptInTags },
+      {
+        onSuccess: () => {
+          invalidateOrganizationsQuery(queryClient)
+          toast.success('Successfully updated AI opt-in settings')
+          onSuccessCallback?.() // Call optional callback on success
+        },
+        onError: (error: ResponseError) => {
+          toast.error(`Failed to update settings: ${error.message}`)
+        },
+      }
+    )
+  }
+
+  return {
+    form,
+    onSubmit,
+    isUpdating,
+    canUpdateOrganization,
+    currentOptInLevel: getAiOptInLevel(selectedOrganization?.opt_in_tags),
+  }
+}

apps/studio/hooks/misc/useOrgOptedIntoAi.ts

@@ -4,6 +4,20 @@ import { OPT_IN_TAGS } from 'lib/constants'
 
 export type AiOptInLevel = 'disabled' | 'schema' | 'schema_and_data'
 
+// Exported helper function
+export const getAiOptInLevel = (tags: string[] | undefined): AiOptInLevel => {
+  const hasSql = tags?.includes(OPT_IN_TAGS.AI_SQL) ?? false
+  const hasData = tags?.includes(OPT_IN_TAGS.AI_DATA ?? 'AI_DATA') // Ensure AI_DATA exists or handle potential undefined
+
+  if (hasSql && hasData) {
+    return 'schema_and_data'
+  } else if (hasSql) {
+    return 'schema'
+  } else {
+    return 'disabled'
+  }
+}
+
 /**
  * Determines if the organization has opted into *any* level of AI features (schema or schema_and_data).
  * This is primarily for backward compatibility.
@@ -36,20 +50,10 @@ export function useOrgOptedIntoAi(): boolean {
  */
 export function useOrgAiOptInLevel(): AiOptInLevel {
   const selectedOrganization = useSelectedOrganization()
-  const optInTags = selectedOrganization?.opt_in_tags ?? []
+  const optInTags = selectedOrganization?.opt_in_tags
 
-  const hasAiSqlTag = optInTags.includes(OPT_IN_TAGS.AI_SQL)
-  const hasAiDataTag = optInTags.includes(OPT_IN_TAGS.AI_DATA)
-
-  let level: AiOptInLevel
-
-  if (hasAiSqlTag && hasAiDataTag) {
-    level = 'schema_and_data'
-  } else if (hasAiSqlTag) {
-    level = 'schema'
-  } else {
-    level = 'disabled'
-  }
+  // Use the helper function
+  const level = getAiOptInLevel(optInTags)
 
   const disallowHipaa = useDisallowHipaa()
   return disallowHipaa(level !== 'disabled') ? level : 'disabled'