Chore/rls ai assistant updates (#19325)

* Have rls ai assistant sit in feature preview + shown via feature flag

* Support editing existing policy in new AI assistant panel

* Beef up feature preview

* Shift AI settings modal to components/ui and add AI settings to RLS AI assistant panel

* Update apps/studio/components/interfaces/App/FeaturePreview/RLSAIAssistantPreview.tsx

Co-authored-by: Copple <10214025+kiwicopple@users.noreply.github.com>

* Update apps/studio/components/interfaces/App/FeaturePreview/RLSAIAssistantPreview.tsx

Co-authored-by: Copple <10214025+kiwicopple@users.noreply.github.com>

---------

Co-authored-by: Copple <10214025+kiwicopple@users.noreply.github.com>

apps/studio/components/interfaces/App/FeaturePreview/FeaturePreviewContext.tsx

@@ -22,6 +22,7 @@ export const FeaturePreviewContextProvider = ({ children }: PropsWithChildren<{}
   const [flags, setFlags] = useState({
     [LOCAL_STORAGE_KEYS.UI_PREVIEW_NAVIGATION_LAYOUT]: false,
     [LOCAL_STORAGE_KEYS.UI_PREVIEW_API_SIDE_PANEL]: false,
+    [LOCAL_STORAGE_KEYS.UI_PREVIEW_RLS_AI_ASSISTANT]: false,
   })
 
   useEffect(() => {
@@ -31,6 +32,8 @@ export const FeaturePreviewContextProvider = ({ children }: PropsWithChildren<{}
           localStorage.getItem(LOCAL_STORAGE_KEYS.UI_PREVIEW_NAVIGATION_LAYOUT) === 'true',
         [LOCAL_STORAGE_KEYS.UI_PREVIEW_API_SIDE_PANEL]:
           localStorage.getItem(LOCAL_STORAGE_KEYS.UI_PREVIEW_API_SIDE_PANEL) === 'true',
+        [LOCAL_STORAGE_KEYS.UI_PREVIEW_RLS_AI_ASSISTANT]:
+          localStorage.getItem(LOCAL_STORAGE_KEYS.UI_PREVIEW_RLS_AI_ASSISTANT) === 'true',
       })
     }
   }, [])
@@ -59,3 +62,8 @@ export const useIsAPIDocsSidePanelEnabled = () => {
   const { flags } = useFeaturePreviewContext()
   return flags[LOCAL_STORAGE_KEYS.UI_PREVIEW_API_SIDE_PANEL]
 }
+
+export const useIsRLSAIAssistantEnabled = () => {
+  const { flags } = useFeaturePreviewContext()
+  return flags[LOCAL_STORAGE_KEYS.UI_PREVIEW_RLS_AI_ASSISTANT]
+}

apps/studio/components/interfaces/App/FeaturePreview/FeaturePreviewModal.tsx

@@ -10,24 +10,32 @@ import Telemetry from 'lib/telemetry'
 import { useAppStateSnapshot } from 'state/app-state'
 import APISidePanelPreview from './APISidePanelPreview'
 import { useFeaturePreviewContext } from './FeaturePreviewContext'
-
-// [Ivan] We should probably move this to a separate file, together with LOCAL_STORAGE_KEYS. We should make adding new feature previews as simple as possible.
-
-const FEATURE_PREVIEWS: { key: string; name: string; content: any; discussionsUrl?: string }[] = [
-  // {
-  //   key: LOCAL_STORAGE_KEYS.UI_PREVIEW_NAVIGATION_LAYOUT,
-  //   name: 'Global navigation update',
-  //   content: null,
-  // },
-  {
-    key: LOCAL_STORAGE_KEYS.UI_PREVIEW_API_SIDE_PANEL,
-    name: 'Project API documentation',
-    content: <APISidePanelPreview />,
-    discussionsUrl: 'https://github.com/orgs/supabase/discussions/18038',
-  },
-]
+import RLSAIAssistantPreview from './RLSAIAssistantPreview'
+import { useFlag } from 'hooks'
 
 const FeaturePreviewModal = () => {
+  const isAiAssistantEnabled = useFlag('policyEditorWithAi')
+
+  // [Ivan] We should probably move this to a separate file, together with LOCAL_STORAGE_KEYS. We should make adding new feature previews as simple as possible.
+  const FEATURE_PREVIEWS: { key: string; name: string; content: any; discussionsUrl?: string }[] = [
+    {
+      key: LOCAL_STORAGE_KEYS.UI_PREVIEW_API_SIDE_PANEL,
+      name: 'Project API documentation',
+      content: <APISidePanelPreview />,
+      discussionsUrl: 'https://github.com/orgs/supabase/discussions/18038',
+    },
+    ...(isAiAssistantEnabled
+      ? [
+          {
+            key: LOCAL_STORAGE_KEYS.UI_PREVIEW_RLS_AI_ASSISTANT,
+            name: 'AI Assistant for RLS policies',
+            content: <RLSAIAssistantPreview />,
+            discussionsUrl: '/', // Need to update
+          },
+        ]
+      : []),
+  ]
+
   const router = useRouter()
   const snap = useAppStateSnapshot()
   const telemetryProps = useTelemetryProps()

apps/studio/components/interfaces/App/FeaturePreview/RLSAIAssistantPreview.tsx

@@ -0,0 +1,50 @@
+import { useParams } from 'common'
+import { Markdown } from 'components/interfaces/Markdown'
+import { BASE_PATH } from 'lib/constants'
+import Image from 'next/image'
+
+const RLSAIAssistantPreview = () => {
+  const { ref } = useParams()
+
+  return (
+    <div className="space-y-2">
+      <div className="mb-4 flex flex-col gap-y-2">
+        <Markdown
+          className="text-foreground-light max-w-full"
+          content={`[Postgres Row Level Security](https://supabase.com/docs/guides/database/postgres/row-level-security) (RLS) is a feature of Postgres that allows you to control which users are allowed to perform operations on specific rows within tables, views, and functions.`}
+        />
+        <Markdown
+          className="text-foreground-light max-w-full"
+          content={`Let our AI Assistant handle the SQL while you focus on building the rules for your policies.`}
+        />
+      </div>
+      <Image
+        src={`${BASE_PATH}/img/previews/rls-ai-assistant-preview.png`}
+        width={1860}
+        height={970}
+        alt="api-docs-side-panel-preview"
+        className="rounded border"
+      />
+      <div className="space-y-2 !mt-4">
+        <p className="text-sm">Enabling this preview will:</p>
+        <ul className="list-disc pl-6 text-sm text-foreground-light space-y-1">
+          <li>
+            <Markdown
+              className="text-foreground-light"
+              content={`Replace the existing wizard-like UI for creating/updating RLS policies [here](/project/${ref}/auth/policies) with a code
+              editor accompanied an AI assistant chat interface`}
+            />
+          </li>
+          <li>
+            <Markdown
+              className="text-foreground-light"
+              content={`Freely explain the scopes of your intended policy, and the AI assistant will generate SQL snippets for you iteratively`}
+            />
+          </li>
+        </ul>
+      </div>
+    </div>
+  )
+}
+
+export default RLSAIAssistantPreview

apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyChat.tsx

@@ -6,16 +6,27 @@ import { useEffect, useMemo, useRef, useState } from 'react'
 import { useForm } from 'react-hook-form'
 import {
   AiIcon,
+  Button,
   FormControl_Shadcn_,
   FormField_Shadcn_,
   FormItem_Shadcn_,
   Form_Shadcn_,
+  IconSettings,
   Input_Shadcn_,
 } from 'ui'
 import * as z from 'zod'
 
 import { useProfile } from 'lib/profile'
 import Message from './Message'
+import { useAppStateSnapshot } from 'state/app-state'
+
+interface AIPolicyChatProps {
+  messages: OpenAI.Beta.Threads.Messages.ThreadMessage[]
+  loading: boolean
+  onSubmit: (s: string) => void
+  onDiff: (s: string) => void
+  onChange: (value: boolean) => void
+}
 
 export const AIPolicyChat = ({
   messages,
@@ -23,15 +34,11 @@ export const AIPolicyChat = ({
   onSubmit,
   onDiff,
   onChange,
-}: {
-  messages: OpenAI.Beta.Threads.Messages.ThreadMessage[]
-  loading: boolean
-  onSubmit: (s: string) => void
-  onDiff: (s: string) => void
-  onChange: (value: boolean) => void
-}) => {
+}: AIPolicyChatProps) => {
   const { profile } = useProfile()
+  const snap = useAppStateSnapshot()
   const bottomRef = useRef<HTMLDivElement>(null)
+
   const name = compact([profile?.first_name, profile?.last_name]).join(' ')
   const sorted = useMemo(() => {
     return sortBy(messages, (m) => m.created_at).filter((m) => {
@@ -82,7 +89,17 @@ export const AIPolicyChat = ({
           }, how can I help you? I'm powered by AI, so surprises and mistakes are possible.
         Make sure to verify any generated code or suggestions, and share feedback so that we can
         learn and improve.`}
-        />
+        >
+          <div>
+            <Button
+              type="default"
+              icon={<IconSettings strokeWidth={1.5} />}
+              onClick={() => snap.setShowAiSettingsModal(true)}
+            >
+              AI Settings
+            </Button>
+          </div>
+        </Message>
 
         {sorted.map((m, idx) => (
           <Message

apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyEditorPanel.utils.ts

@@ -1,3 +1,4 @@
+import { PostgresPolicy } from '@supabase/postgres-meta'
 import { uuidv4 } from 'lib/helpers'
 import { ThreadMessage } from 'openai/resources/beta/threads/messages/messages'
 
@@ -33,3 +34,43 @@ export const generateThreadMessage = ({
   }
   return message
 }
+
+export const generatePlaceholder = (policy?: PostgresPolicy) => {
+  if (policy === undefined) {
+    return `
+CREATE POLICY *name* ON *table_name*\n
+AS PERMISSIVE -- PERMISSIVE | RESTRICTIVE\n
+FOR ALL -- ALL | SELECT | INSERT | UPDATE | DELETE\n
+TO *role_name* -- Default: public\n
+USING ( *using_expression* )\n
+WITH CHECK ( *check_expression* );
+`.trim()
+  } else {
+    let expression = ''
+    if (policy.definition !== null && policy.definition !== undefined) {
+      expression += `USING ( *${policy.definition}* )${
+        policy.check === null || policy.check === undefined ? ';' : ''
+      }\n`
+    }
+    if (policy.check !== null && policy.check !== undefined) {
+      expression += `WITH CHECK ( *${policy.check}* );\n`
+    }
+
+    return `
+BEGIN;\n
+ \n
+-- To update your policy definition\n
+ALTER POLICY "${policy.name}"\n
+ON "${policy.schema}"."${policy.table}"\n
+TO *${policy.roles.join(', ')}*\n
+${expression}
+ \n
+-- To rename your policy\n
+ALTER POLICY "${policy.name}"\n
+ON "${policy.schema}"."${policy.table}"\n
+RENAME TO "*New Policy Name*";\n
+ \n
+COMMIT;
+`.trim()
+  }
+}

apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/AIPolicyHeader.tsx

@@ -1,24 +1,37 @@
+import { PostgresPolicy } from '@supabase/postgres-meta'
+import styles from '@ui/layout/ai-icon-animation/ai-icon-animation-style.module.css'
+import clsx from 'clsx'
+import { X } from 'lucide-react'
 import { AiIcon, Button, SheetClose_Shadcn_, SheetHeader_Shadcn_, SheetTitle_Shadcn_, cn } from 'ui'
 
-import styles from '@ui/layout/ai-icon-animation/ai-icon-animation-style.module.css'
-import { X } from 'lucide-react'
-
 export const AIPolicyHeader = ({
+  selectedPolicy,
   assistantVisible,
   setAssistantVisible,
 }: {
+  selectedPolicy?: PostgresPolicy
   assistantVisible: boolean
   setAssistantVisible: (v: boolean) => void
 }) => {
   return (
     <SheetHeader_Shadcn_ className="py-3 flex flex-row justify-between items-center">
-      <div className="flex flex-row gap-3 items-center">
-        <SheetClose_Shadcn_ className="text-light hover:text ring-offset-background transition-opacity hover:opacity-100 focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2 disabled:pointer-events-none data-[state=open]:bg-secondary">
+      <div className="flex flex-row gap-3 items-center max-w-[75%]">
+        <SheetClose_Shadcn_
+          className={clsx(
+            'text-light hover:text ring-offset-background transition-opacity hover:opacity-100',
+            'focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2',
+            'disabled:pointer-events-none data-[state=open]:bg-secondary'
+          )}
+        >
           <X className="h-3 w-3" />
           <span className="sr-only">Close</span>
         </SheetClose_Shadcn_>
-        <div className="h-[24px] w-[1px] bg-border"></div>
-        <SheetTitle_Shadcn_>Create a new row level security policy</SheetTitle_Shadcn_>
+        <div className="h-[24px] w-[1px] bg-border" />
+        <SheetTitle_Shadcn_ className="truncate">
+          {selectedPolicy !== undefined
+            ? `Update policy: ${selectedPolicy.name}`
+            : 'Create a new row level security policy'}
+        </SheetTitle_Shadcn_>
       </div>
       <Button
         aria-expanded={assistantVisible}

apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/Message.tsx

@@ -3,14 +3,23 @@ import dayjs from 'dayjs'
 import { kebabCase, noop, take } from 'lodash'
 import { Copy, FileDiff } from 'lucide-react'
 import Image from 'next/image'
-import { memo, useMemo } from 'react'
+import { PropsWithChildren, memo, useMemo } from 'react'
 import ReactMarkdown from 'react-markdown'
 import { format } from 'sql-formatter'
-import { AiIcon, AiIconAnimation, Badge, Button } from 'ui'
+import { AiIconAnimation, Badge, Button, IconSettings } from 'ui'
 
 import CodeEditor from 'components/ui/CodeEditor'
 import { useProfile } from 'lib/profile'
 
+interface MessageProps {
+  name?: string
+  role: 'user' | 'assistant'
+  content?: string
+  createdAt?: number
+  isDebug?: boolean
+  onDiff?: (s: string) => void
+}
+
 const Message = memo(function Message({
   name,
   role,
@@ -18,14 +27,8 @@ const Message = memo(function Message({
   createdAt,
   isDebug,
   onDiff = noop,
-}: {
-  name?: string
-  role: 'user' | 'assistant'
-  content?: string
-  createdAt?: number
-  isDebug?: boolean
-  onDiff?: (s: string) => void
-}) {
+  children,
+}: PropsWithChildren<MessageProps>) {
   const { profile } = useProfile()
 
   const icon = useMemo(() => {
@@ -145,6 +148,7 @@ const Message = memo(function Message({
       >
         {content}
       </ReactMarkdown>
+      {children}
     </div>
   )
 })

apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/RLSCodeEditor.tsx

@@ -16,35 +16,17 @@ interface RLSCodeEditorProps {
   wrapperClassName?: string
   className?: string
   value?: string
+  placeholder?: string
   editorRef: MutableRefObject<editor.IStandaloneCodeEditor | null>
 }
 
-// const placeholderText = `
-// CREATE POLICY *name* ON *table_name*\n
-// [ AS { PERMISSIVE | RESTRICTIVE } ]\n
-// [ FOR { ALL | SELECT | INSERT | UPDATE | DELETE } ]\n
-// [ TO *role_name* ]\n
-// [ USING ( *using_expression* ) ]\n
-// [ WITH CHECK ( *check_expression* ) ];
-// `.trim()
-
-const placeholderText = `
-CREATE POLICY *name* ON *table_name*\n
-AS PERMISSIVE -- PERMISSIVE | RESTRICTIVE\n
-FOR ALL -- ALL | SELECT | INSERT | UPDATE | DELETE\n
-TO *role_name* -- Default: public\n
-USING ( *using_expression* )\n
-WITH CHECK ( *check_expression* );\n
-&nbsp;\n
--- Docs: https://www.postgresql.org/docs/current/sql-createpolicy.html
-`.trim()
-
 const RLSCodeEditor = ({
   id,
   defaultValue,
   wrapperClassName,
   className,
   value,
+  placeholder,
   editorRef,
 }: RLSCodeEditorProps) => {
   const hasValue = useRef<any>()
@@ -55,8 +37,8 @@ const RLSCodeEditor = ({
 
     hasValue.current = editor.createContextKey('hasValue', false)
 
-    const placeholder = document.querySelector('.monaco-placeholder') as HTMLElement | null
-    if (placeholder) placeholder.style.display = 'block'
+    const placeholderEl = document.querySelector('.monaco-placeholder') as HTMLElement | null
+    if (placeholderEl) placeholderEl.style.display = 'block'
 
     editor.addCommand(
       monaco.KeyCode.Tab,
@@ -66,7 +48,7 @@ const RLSCodeEditor = ({
             // @ts-ignore
             identifier: 'add-placeholder',
             range: new monaco.Range(1, 1, 1, 1),
-            text: placeholderText
+            text: (placeholder ?? '')
               .split('\n\n')
               .join('\n')
               .replaceAll('*', '')
@@ -83,12 +65,12 @@ const RLSCodeEditor = ({
   const onChange: OnChange = (value) => {
     hasValue.current.set((value ?? '').length > 0)
 
-    const placeholder = document.querySelector('.monaco-placeholder') as HTMLElement | null
-    if (placeholder) {
+    const placeholderEl = document.querySelector('.monaco-placeholder') as HTMLElement | null
+    if (placeholderEl) {
       if (!value) {
-        placeholder.style.display = 'block'
+        placeholderEl.style.display = 'block'
       } else {
-        placeholder.style.display = 'none'
+        placeholderEl.style.display = 'none'
       }
     }
   }
@@ -122,12 +104,14 @@ const RLSCodeEditor = ({
         onMount={onMount}
         onChange={onChange}
       />
-      <div
-        className="monaco-placeholder absolute top-[3px] left-[57px] text-sm pointer-events-none font-mono [&>div>p]:text-foreground-lighter [&>div>p]:!m-0 tracking-tighter"
-        style={{ display: 'none' }}
-      >
-        <Markdown content={placeholderText} />
-      </div>
+      {placeholder !== undefined && (
+        <div
+          className="monaco-placeholder absolute top-[3px] left-[57px] text-sm pointer-events-none font-mono [&>div>p]:text-foreground-lighter [&>div>p]:!m-0 tracking-tighter"
+          style={{ display: 'none' }}
+        >
+          <Markdown content={placeholder} />
+        </div>
+      )}
     </>
   )
 }

apps/studio/components/interfaces/Auth/Policies/AIPolicyEditorPanel/index.tsx

@@ -1,3 +1,4 @@
+import { PostgresPolicy } from '@supabase/postgres-meta'
 import { FileDiff } from 'lucide-react'
 import dynamic from 'next/dynamic'
 import { ThreadMessage } from 'openai/resources/beta/threads/messages/messages'
@@ -18,7 +19,7 @@ import { QueryResponseError, useExecuteSqlMutation } from 'data/sql/execute-sql-
 import { useSelectedProject, useStore } from 'hooks'
 import { uuidv4 } from 'lib/helpers'
 import { AIPolicyChat } from './AIPolicyChat'
-import { generateThreadMessage } from './AIPolicyEditorPanel.utils'
+import { generatePlaceholder, generateThreadMessage } from './AIPolicyEditorPanel.utils'
 import { AIPolicyHeader } from './AIPolicyHeader'
 import QueryError from './QueryError'
 import RLSCodeEditor from './RLSCodeEditor'
@@ -30,6 +31,7 @@ const DiffEditor = dynamic(
 
 interface AIPolicyEditorPanelProps {
   visible: boolean
+  selectedPolicy?: PostgresPolicy
   onSelectCancel: () => void
 }
 
@@ -38,6 +40,7 @@ interface AIPolicyEditorPanelProps {
  */
 export const AIPolicyEditorPanel = memo(function ({
   visible,
+  selectedPolicy,
   onSelectCancel,
 }: AIPolicyEditorPanelProps) {
   const { meta } = useStore()
@@ -45,6 +48,7 @@ export const AIPolicyEditorPanel = memo(function ({
 
   const editorRef = useRef<IStandaloneCodeEditor | null>(null)
   const diffEditorRef = useRef<IStandaloneDiffEditor | null>(null)
+  const placeholder = generatePlaceholder(selectedPolicy)
 
   const [error, setError] = useState<QueryResponseError>()
   // [Joshen] Separate state here as there's a delay between submitting and the API updating the loading status
@@ -142,9 +146,9 @@ export const AIPolicyEditorPanel = memo(function ({
   const errorLines =
     error?.formattedError.split('\n').filter((x: string) => x.length > 0).length ?? 0
 
-  const createNewPolicy = useCallback(() => {
+  const onExecuteSQL = useCallback(() => {
     // clean up the sql before sending
-    const policy = editorRef.current?.getValue().replaceAll('\n', ' ').replaceAll('  ', ' ')
+    const policy = editorRef.current?.getValue().replaceAll('  ', ' ')
 
     if (policy) {
       setError(undefined)
@@ -255,10 +259,14 @@ export const AIPolicyEditorPanel = memo(function ({
       <Sheet_Shadcn_ open={visible} onOpenChange={() => onClosingPanel()}>
         <SheetContent_Shadcn_
           size={assistantVisible ? 'lg' : 'default'}
-          className={cn('p-0 flex flex-row gap-0', assistantVisible && '!min-w-[1024px]')}
+          className={cn(
+            'p-0 flex flex-row gap-0',
+            assistantVisible ? '!min-w-[1024px]' : '!min-w-[600px]'
+          )}
         >
           <div className={cn('flex flex-col grow w-full', assistantVisible && 'w-[60%]')}>
             <AIPolicyHeader
+              selectedPolicy={selectedPolicy}
               assistantVisible={assistantVisible}
               setAssistantVisible={setAssistantPanel}
             />
@@ -311,6 +319,7 @@ export const AIPolicyEditorPanel = memo(function ({
                   wrapperClassName={incomingChange ? '!hidden' : ''}
                   defaultValue={''}
                   editorRef={editorRef}
+                  placeholder={placeholder}
                 />
               </div>
 
@@ -325,7 +334,7 @@ export const AIPolicyEditorPanel = memo(function ({
                       loading={isExecuting}
                       htmlType="submit"
                       disabled={isExecuting || incomingChange !== undefined}
-                      onClick={() => createNewPolicy()}
+                      onClick={() => onExecuteSQL()}
                     >
                       Save policy
                     </Button>

apps/studio/components/interfaces/Auth/Policies/Policies.tsx

@@ -14,19 +14,27 @@ import ProductEmptyState from 'components/to-be-cleaned/ProductEmptyState'
 import ConfirmModal from 'components/ui/Dialogs/ConfirmDialog'
 import InformationBox from 'components/ui/InformationBox'
 import { tableKeys } from 'data/tables/keys'
+import { useIsRLSAIAssistantEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext'
 
 interface PoliciesProps {
   tables: PostgresTable[]
   hasTables: boolean
   isLocked: boolean
+  onSelectEditPolicy: (policy: PostgresPolicy) => void
 }
 
-const Policies = ({ tables, hasTables, isLocked }: PoliciesProps) => {
+const Policies = ({
+  tables,
+  hasTables,
+  isLocked,
+  onSelectEditPolicy: onSelectEditPolicyAI,
+}: PoliciesProps) => {
   const router = useRouter()
   const { ref } = useParams()
 
   const { ui, meta } = useStore()
   const queryClient = useQueryClient()
+  const isAiAssistantEnabled = useIsRLSAIAssistantEnabled()
   const roles = meta.roles.list((role: PostgresRole) => !meta.roles.systemRoles.includes(role.name))
 
   const [selectedSchemaAndTable, setSelectedSchemaAndTable] = useState<any>({})
@@ -57,8 +65,12 @@ const Policies = ({ tables, hasTables, isLocked }: PoliciesProps) => {
   }
 
   const onSelectEditPolicy = (policy: any) => {
-    setSelectedPolicyToEdit(policy)
-    setSelectedSchemaAndTable({ schema: policy.schema, table: policy.table })
+    if (isAiAssistantEnabled) {
+      onSelectEditPolicyAI(policy)
+    } else {
+      setSelectedPolicyToEdit(policy)
+      setSelectedSchemaAndTable({ schema: policy.schema, table: policy.table })
+    }
   }
 
   const onSelectDeletePolicy = (policy: any) => {

apps/studio/components/interfaces/Auth/Policies/PolicyTableRow/PolicyTableRowHeader.tsx

@@ -7,6 +7,7 @@ import { useRouter } from 'next/router'
 import { Badge, Button, IconLock } from 'ui'
 
 import { useCheckPermissions } from 'hooks'
+import { useIsRLSAIAssistantEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext'
 
 interface PolicyTableRowHeaderProps {
   table: PostgresTable
@@ -23,6 +24,7 @@ const PolicyTableRowHeader = ({
 }: PolicyTableRowHeaderProps) => {
   const router = useRouter()
   const { ref } = router.query
+  const isAiAssistantEnabled = useIsRLSAIAssistantEnabled()
   const canToggleRLS = useCheckPermissions(PermissionAction.TENANT_SQL_ADMIN_WRITE, 'tables')
   const canCreatePolicies = useCheckPermissions(PermissionAction.TENANT_SQL_ADMIN_WRITE, 'policies')
 
@@ -46,36 +48,7 @@ const PolicyTableRowHeader = ({
       </div>
       {!isLocked && (
         <div className="flex-1">
-          <div className="flex flex-row-reverse">
-            <Tooltip.Root delayDuration={0}>
-              <Tooltip.Trigger>
-                <Button
-                  type="outline"
-                  disabled={!canCreatePolicies}
-                  className="ml-2"
-                  onClick={() => onSelectCreatePolicy(table)}
-                >
-                  New Policy
-                </Button>
-              </Tooltip.Trigger>
-              {!canCreatePolicies && (
-                <Tooltip.Portal>
-                  <Tooltip.Content side="bottom">
-                    <Tooltip.Arrow className="radix-tooltip-arrow" />
-                    <div
-                      className={[
-                        'rounded bg-alternative py-1 px-2 leading-none shadow',
-                        'border border-background',
-                      ].join(' ')}
-                    >
-                      <span className="text-xs text-foreground">
-                        You need additional permissions to create RLS policies
-                      </span>
-                    </div>
-                  </Tooltip.Content>
-                </Tooltip.Portal>
-              )}
-            </Tooltip.Root>
+          <div className="flex flex-row justify-end gap-x-2">
             <Tooltip.Root delayDuration={0}>
               <Tooltip.Trigger>
                 <Button
@@ -104,6 +77,36 @@ const PolicyTableRowHeader = ({
                 </Tooltip.Portal>
               )}
             </Tooltip.Root>
+            {!isAiAssistantEnabled && (
+              <Tooltip.Root delayDuration={0}>
+                <Tooltip.Trigger>
+                  <Button
+                    type="outline"
+                    disabled={!canCreatePolicies}
+                    onClick={() => onSelectCreatePolicy(table)}
+                  >
+                    New Policy
+                  </Button>
+                </Tooltip.Trigger>
+                {!canCreatePolicies && (
+                  <Tooltip.Portal>
+                    <Tooltip.Content side="bottom">
+                      <Tooltip.Arrow className="radix-tooltip-arrow" />
+                      <div
+                        className={[
+                          'rounded bg-alternative py-1 px-2 leading-none shadow',
+                          'border border-background',
+                        ].join(' ')}
+                      >
+                        <span className="text-xs text-foreground">
+                          You need additional permissions to create RLS policies
+                        </span>
+                      </div>
+                    </Tooltip.Content>
+                  </Tooltip.Portal>
+                )}
+              </Tooltip.Root>
+            )}
           </div>
         </div>
       )}

apps/studio/components/interfaces/SQLEditor/SQLEditor.tsx

@@ -31,6 +31,7 @@ import { SqlSnippet } from 'data/content/sql-snippets-query'
 import { useEntityDefinitionsQuery } from 'data/database/entity-definitions-query'
 import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
 import { useFormatQueryMutation } from 'data/sql/format-sql-query'
+import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query'
 import { isError } from 'data/utils/error-check'
 import {
   useFlag,
@@ -44,10 +45,10 @@ import { IS_PLATFORM, OPT_IN_TAGS } from 'lib/constants'
 import { uuidv4 } from 'lib/helpers'
 import { useProfile } from 'lib/profile'
 import Telemetry from 'lib/telemetry'
+import { useAppStateSnapshot } from 'state/app-state'
 import { getSqlEditorStateSnapshot, useSqlEditorStateSnapshot } from 'state/sql-editor'
 import { subscriptionHasHipaaAddon } from '../Billing/Subscription/Subscription.utils'
 import AISchemaSuggestionPopover from './AISchemaSuggestionPopover'
-import AISettingsModal from './AISettingsModal'
 import { sqlAiDisclaimerComment, untitledSnippetTitle } from './SQLEditor.constants'
 import {
   ContentDiff,
@@ -63,7 +64,6 @@ import {
   getDiffTypeDropdownLabel,
 } from './SQLEditor.utils'
 import UtilityPanel from './UtilityPanel/UtilityPanel'
-import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query'
 
 // Load the monaco editor client-side only (does not behave well server-side)
 const MonacoEditor = dynamic(() => import('./MonacoEditor'), { ssr: false })
@@ -99,6 +99,7 @@ const SQLEditor = () => {
   const { profile } = useProfile()
   const project = useSelectedProject()
   const organization = useSelectedOrganization()
+  const appSnap = useAppStateSnapshot()
   const snap = useSqlEditorStateSnapshot()
 
   const { mutate: formatQuery } = useFormatQueryMutation()
@@ -121,8 +122,6 @@ const SQLEditor = () => {
   const hasHipaaAddon = subscriptionHasHipaaAddon(subscription)
 
   const [isAiOpen, setIsAiOpen] = useLocalStorageQuery('supabase_sql-editor-ai-open', true)
-
-  const [isAISettingsOpen, setIsAISettingsOpen] = useState(false)
   const [isConfirmModalOpen, setIsConfirmModalOpen] = useState(false)
 
   const selectedOrganization = useSelectedOrganization()
@@ -514,7 +513,6 @@ const SQLEditor = () => {
         setDebugSolution,
       }}
     >
-      <AISettingsModal visible={isAISettingsOpen} onCancel={() => setIsAISettingsOpen(false)} />
       <ConfirmModal
         visible={isConfirmModalOpen}
         title="Destructive operation"
@@ -533,7 +531,7 @@ const SQLEditor = () => {
         {isAiOpen && supabaseAIEnabled && !hasHipaaAddon && (
           <AISchemaSuggestionPopover
             onClickSettings={() => {
-              setIsAISettingsOpen(true)
+              appSnap.setShowAiSettingsModal(true)
             }}
           >
             <motion.div
@@ -797,7 +795,7 @@ const SQLEditor = () => {
                       <button
                         onClick={() => {
                           setIsSchemaSuggestionDismissed(true)
-                          setIsAISettingsOpen(true)
+                          appSnap.setShowAiSettingsModal(true)
                         }}
                         className="text-brand-600 hover:text-brand-600 transition"
                       >

apps/studio/components/layouts/ProjectLayout/ProjectLayout.tsx

@@ -1,4 +1,7 @@
 import { useParams } from 'common/hooks'
+import ProjectAPIDocs from 'components/interfaces/ProjectAPIDocs/ProjectAPIDocs'
+import AISettingsModal from 'components/ui/AISettingsModal'
+import Connecting from 'components/ui/Loading/Loading'
 import ResourceExhaustionWarningBanner from 'components/ui/ResourceExhaustionWarningBanner/ResourceExhaustionWarningBanner'
 import { useFlag, useSelectedOrganization, useSelectedProject, withAuth } from 'hooks'
 import { IS_PLATFORM, PROJECT_STATUS } from 'lib/constants'
@@ -18,8 +21,6 @@ import { ProjectContextProvider } from './ProjectContext'
 import ProjectPausedState from './ProjectPausedState'
 import RestoringState from './RestoringState'
 import UpgradingState from './UpgradingState'
-import Connecting from 'components/ui/Loading/Loading'
-import ProjectAPIDocs from 'components/interfaces/ProjectAPIDocs/ProjectAPIDocs'
 
 // [Joshen] This is temporary while we unblock users from managing their project
 // if their project is not responding well for any reason. Eventually needs a bit of an overhaul
@@ -120,6 +121,7 @@ const ProjectLayout = ({
         </div>
 
         <EnableBranchingModal />
+        <AISettingsModal />
         <ProjectAPIDocs />
       </ProjectContextProvider>
     </AppLayout>
@@ -269,6 +271,7 @@ export const ProjectLayoutNonBlocking = ({
         </div>
 
         <EnableBranchingModal />
+        <AISettingsModal />
         <ProjectAPIDocs />
       </ProjectContextProvider>
     </AppLayout>

apps/studio/components/ui/AISettingsModal.tsx

@@ -1,13 +1,13 @@
-import { ModalProps } from '@ui/components/Modal/Modal'
-import { useLocalStorageQuery, useSelectedOrganization, useStore } from 'hooks'
-import { IS_PLATFORM, OPT_IN_TAGS } from 'lib/constants'
 import Link from 'next/link'
 import { Alert, IconExternalLink, Modal, Toggle } from 'ui'
-export interface AISettingsModalProps {
-  visible: boolean
-}
 
-const AISettingsModal = (props: ModalProps) => {
+import { useLocalStorageQuery, useSelectedOrganization, useStore } from 'hooks'
+import { IS_PLATFORM, OPT_IN_TAGS } from 'lib/constants'
+import { useAppStateSnapshot } from 'state/app-state'
+
+const AISettingsModal = () => {
+  const snap = useAppStateSnapshot()
+
   const selectedOrganization = useSelectedOrganization()
   const isOptedInToAI = selectedOrganization?.opt_in_tags?.includes(OPT_IN_TAGS.AI_SQL) ?? false
   const [hasEnabledAISchema, setHasEnabledAISchema] = useLocalStorageQuery(
@@ -24,7 +24,12 @@ const AISettingsModal = (props: ModalProps) => {
   }
 
   return (
-    <Modal header="SQL Editor AI Settings" hideFooter closable {...props}>
+    <Modal
+      header="Supabase AI Settings"
+      hideFooter
+      visible={snap.showAiSettingsModal}
+      onCancel={() => snap.setShowAiSettingsModal(false)}
+    >
       <div className="flex flex-col items-start justify-between gap-4 px-6 py-3">
         {IS_PLATFORM && !isOptedInToAI && selectedOrganization && (
           <Alert

apps/studio/lib/constants/index.ts

@@ -98,6 +98,7 @@ export const LOCAL_STORAGE_KEYS = {
   TELEMETRY_CONSENT: 'supabase-consent',
   UI_PREVIEW_NAVIGATION_LAYOUT: 'supabase-ui-preview-nav-layout',
   UI_PREVIEW_API_SIDE_PANEL: 'supabase-ui-api-side-panel',
+  UI_PREVIEW_RLS_AI_ASSISTANT: 'supabase-ui-rls-ai-assistant',
   DASHBOARD_HISTORY: (ref: string) => `dashboard-history-${ref}`,
 }
 

apps/studio/pages/project/[ref]/auth/policies.tsx

@@ -7,6 +7,7 @@ import { observer } from 'mobx-react-lite'
 import { useEffect, useState } from 'react'
 import { Button, IconExternalLink, IconSearch, Input } from 'ui'
 
+import { useIsRLSAIAssistantEnabled } from 'components/interfaces/App/FeaturePreview/FeaturePreviewContext'
 import { Policies } from 'components/interfaces/Auth/Policies'
 import { AIPolicyEditorPanel } from 'components/interfaces/Auth/Policies/AIPolicyEditorPanel'
 import { AuthLayout } from 'components/layouts'
@@ -17,7 +18,7 @@ import SchemaSelector from 'components/ui/SchemaSelector'
 import { GenericSkeletonLoader } from 'components/ui/ShimmeringLoader'
 import { useSchemasQuery } from 'data/database/schemas-query'
 import { useTablesQuery } from 'data/tables/tables-query'
-import { useCheckPermissions, useFlag, useStore } from 'hooks'
+import { useCheckPermissions, useStore } from 'hooks'
 import { EXCLUDED_SCHEMAS } from 'lib/constants/schemas'
 import { useTableEditorStateSnapshot } from 'state/table-editor'
 import { NextPageWithLayout } from 'types'
@@ -63,9 +64,10 @@ const AuthPoliciesPage: NextPageWithLayout = () => {
   const { meta } = useStore()
   const { search } = useParams()
   const [searchString, setSearchString] = useState<string>('')
-  const canCreatePolicyWithAi = useFlag('policyEditorWithAi')
 
   const [showPolicyAiEditor, setShowPolicyAiEditor] = useState(false)
+  const [selectedPolicyToEdit, setSelectedPolicyToEdit] = useState<PostgresPolicy>()
+  const isAiAssistantEnabled = useIsRLSAIAssistantEnabled()
 
   useEffect(() => {
     if (search) setSearchString(search)
@@ -137,7 +139,7 @@ const AuthPoliciesPage: NextPageWithLayout = () => {
                 Documentation
               </Button>
             </a>
-            {canCreatePolicyWithAi && (
+            {isAiAssistantEnabled && (
               <Tooltip.Root delayDuration={0}>
                 <Tooltip.Trigger>
                   <Button
@@ -176,12 +178,24 @@ const AuthPoliciesPage: NextPageWithLayout = () => {
       {isError && <AlertError error={error} subject="Failed to retrieve tables" />}
 
       {isSuccess && (
-        <Policies tables={filteredTables} hasTables={tables.length > 0} isLocked={isLocked} />
+        <Policies
+          tables={filteredTables}
+          hasTables={tables.length > 0}
+          isLocked={isLocked}
+          onSelectEditPolicy={(policy) => {
+            setSelectedPolicyToEdit(policy)
+            setShowPolicyAiEditor(true)
+          }}
+        />
       )}
 
       <AIPolicyEditorPanel
         visible={showPolicyAiEditor}
-        onSelectCancel={() => setShowPolicyAiEditor(false)}
+        selectedPolicy={selectedPolicyToEdit}
+        onSelectCancel={() => {
+          setShowPolicyAiEditor(false)
+          setSelectedPolicyToEdit(undefined)
+        }}
       />
     </div>
   )

apps/studio/state/app-state.ts

@@ -50,6 +50,10 @@ export const appState = proxy({
   setShowFeaturePreviewModal: (value: boolean) => {
     appState.showFeaturePreviewModal = value
   },
+  showAiSettingsModal: false,
+  setShowAiSettingsModal: (value: boolean) => {
+    appState.showAiSettingsModal = value
+  },
 })
 
 export const getAppStateSnapshot = () => snapshot(appState)