Review branch diff with AI (#36915)

apps/studio/components/interfaces/BranchManagement/ReviewWithAI.tsx

@@ -0,0 +1,123 @@
+import { AiIconAnimation } from 'ui'
+import { useProjectByRef } from 'hooks/misc/useSelectedProject'
+import { useTablesQuery } from 'data/tables/tables-query'
+import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state'
+import { Branch } from 'data/branches/branches-query'
+import { tablesToSQL } from 'lib/helpers'
+import { ButtonTooltip } from 'components/ui/ButtonTooltip'
+
+interface ReviewWithAIProps {
+  currentBranch?: Branch
+  mainBranch?: Branch
+  parentProjectRef?: string
+  diffContent?: string
+  disabled?: boolean
+}
+
+export const ReviewWithAI = ({
+  currentBranch,
+  mainBranch,
+  parentProjectRef,
+  diffContent,
+  disabled = false,
+}: ReviewWithAIProps) => {
+  const aiSnap = useAiAssistantStateSnapshot()
+
+  // Get parent project for production schema
+  const parentProject = useProjectByRef(parentProjectRef)
+
+  // Fetch production schema tables
+  const { data: productionTables } = useTablesQuery(
+    {
+      projectRef: parentProjectRef,
+      connectionString: (parentProject as any)?.connectionString,
+      schema: 'public',
+      includeColumns: true,
+    },
+    { enabled: !!parentProjectRef && !!parentProject }
+  )
+
+  const handleReviewWithAssistant = () => {
+    if (!currentBranch || !mainBranch) return
+
+    // Prepare diff content for the assistant
+    const sqlSnippets = []
+
+    // Add production schema SQL if available
+    if (productionTables && productionTables.length > 0) {
+      const productionSQL = tablesToSQL(productionTables)
+      if (productionSQL.trim()) {
+        sqlSnippets.push({
+          label: 'Production Schema',
+          content: productionSQL,
+        })
+      }
+    }
+
+    // Add database diff content if available
+    if (diffContent && diffContent.trim()) {
+      sqlSnippets.push({
+        label: 'Database Changes',
+        content: '-- DATABASE CHANGES:\n' + diffContent,
+      })
+    }
+
+    aiSnap.newChat({
+      name: `Review merge: ${currentBranch.name} → ${mainBranch.name}`,
+      open: true,
+      sqlSnippets: sqlSnippets.length > 0 ? sqlSnippets : undefined,
+      initialInput: `Please review this merge request from branch "${currentBranch.name}" into "${mainBranch.name || 'main'}". 
+
+I've included the current production schema as context, along with the proposed database changes.
+
+Analyze the changes and provide feedback on:
+- Database schema changes and potential impacts on the production schema
+- Migration safety and rollback considerations
+- Overall code quality and best practices
+- Potential breaking changes or compatibility issues
+- Data integrity and constraint implications
+
+Please be concise with your response.`,
+      suggestions: {
+        title: `I can help you review the database schema changes from "${currentBranch.name}" to "${mainBranch.name}", here are some specific areas I can focus on:`,
+        prompts: [
+          {
+            label: 'Schema Impact',
+            description:
+              'Analyze the database schema changes and their potential impact on production...',
+          },
+          {
+            label: 'Migration Safety',
+            description: 'Review the migration safety and rollback strategies...',
+          },
+          {
+            label: 'Performance',
+            description: 'Analyze potential performance implications of these changes...',
+          },
+          {
+            label: 'Data Integrity',
+            description: 'Review constraints, indexes, and data integrity implications...',
+          },
+        ],
+      },
+    })
+  }
+
+  return (
+    <ButtonTooltip
+      type="default"
+      disabled={disabled || !currentBranch || !mainBranch}
+      className="px-1"
+      onClick={handleReviewWithAssistant}
+      tooltip={{
+        content: {
+          side: 'bottom',
+          text: 'Ask Supabase Assistant to review the merge request',
+        },
+      }}
+    >
+      <AiIconAnimation size={16} />
+      <span className="sr-only">Review with Assistant</span>
+    </ButtonTooltip>
+  )
+}

apps/studio/components/interfaces/Database/Schemas/SchemaGraph.tsx

@@ -22,6 +22,7 @@ import { SchemaGraphLegend } from './SchemaGraphLegend'
 import { getGraphDataFromTables, getLayoutedElementsViaDagre } from './Schemas.utils'
 import { TableNode } from './SchemaTableNode'
 import { copyToClipboard } from 'ui'
+import { tablesToSQL } from 'lib/helpers'
 // [Joshen] Persisting logic: Only save positions to local storage WHEN a node is moved OR when explicitly clicked to reset layout
 
 export const SchemaGraph = () => {
@@ -167,59 +168,6 @@ export const SchemaGraph = () => {
     }
   }
 
-  function tablesToSQL(t: typeof tables) {
-    if (!Array.isArray(t)) return ''
-    const warning =
-      '-- WARNING: This schema is for context only and is not meant to be run.\n-- Table order and constraints may not be valid for execution.\n\n'
-    const sql = t
-      .map((table) => {
-        if (!table || !Array.isArray((table as any).columns)) return ''
-
-        const columns = (table as { columns?: any[] }).columns ?? []
-        const columnLines = columns.map((c) => {
-          let line = `  ${c.name} ${c.data_type}`
-          if (c.is_identity) {
-            line += ' GENERATED ALWAYS AS IDENTITY'
-          }
-          if (c.is_nullable === false) {
-            line += ' NOT NULL'
-          }
-          if (c.default_value !== null && c.default_value !== undefined) {
-            line += ` DEFAULT ${c.default_value}`
-          }
-          if (c.is_unique) {
-            line += ' UNIQUE'
-          }
-          if (c.check) {
-            line += ` CHECK (${c.check})`
-          }
-          return line
-        })
-
-        const constraints: string[] = []
-
-        if (Array.isArray(table.primary_keys) && table.primary_keys.length > 0) {
-          const pkCols = table.primary_keys.map((pk) => pk.name).join(', ')
-          constraints.push(`  CONSTRAINT ${table.name}_pkey PRIMARY KEY (${pkCols})`)
-        }
-
-        if (Array.isArray(table.relationships)) {
-          table.relationships.forEach((rel) => {
-            if (rel && rel.source_table_name === table.name) {
-              constraints.push(
-                `  CONSTRAINT ${rel.constraint_name} FOREIGN KEY (${rel.source_column_name}) REFERENCES ${rel.target_table_schema}.${rel.target_table_name}(${rel.target_column_name})`
-              )
-            }
-          })
-        }
-
-        const allLines = [...columnLines, ...constraints]
-        return `CREATE TABLE ${table.schema}.${table.name} (\n${allLines.join(',\n')}\n);`
-      })
-      .join('\n')
-    return warning + sql
-  }
-
   useEffect(() => {
     if (isSuccessTables && isSuccessSchemas && tables.length > 0) {
       const schema = schemas.find((s) => s.name === selectedSchema) as PostgresSchema

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

@@ -30,6 +30,7 @@ import { AIAssistantChatSelector } from './AIAssistantChatSelector'
 import { AIOnboarding } from './AIOnboarding'
 import { AIOptInModal } from './AIOptInModal'
 import { AssistantChatForm } from './AssistantChatForm'
+import { type SqlSnippet } from './AIAssistant.types'
 import { Message } from './Message'
 import { useAutoScroll } from './hooks'
 import type { AssistantMessageType } from 'state/ai-assistant-state'
@@ -305,7 +306,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
 
   useEffect(() => {
     if (snap.open && isInSQLEditor && !!snippetContent) {
-      snap.setSqlSnippets([snippetContent])
+      snap.setSqlSnippets([{ label: 'Current Query', content: snippetContent }])
     }
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [snap.open, isInSQLEditor, snippetContent])
@@ -483,7 +484,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
               onMessageSend={sendMessageToAssistant}
               value={value}
               onValueChange={setValue}
-              sqlSnippets={snap.sqlSnippets as string[] | undefined}
+              sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
               onRemoveSnippet={(index) => {
                 const newSnippets = [...(snap.sqlSnippets ?? [])]
                 newSnippets.splice(index, 1)
@@ -581,7 +582,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
                 sendMessageToAssistant(finalMessage)
                 scrollToEnd()
               }}
-              sqlSnippets={snap.sqlSnippets as string[] | undefined}
+              sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
               onRemoveSnippet={(index) => {
                 const newSnippets = [...(snap.sqlSnippets ?? [])]
                 newSnippets.splice(index, 1)

apps/studio/components/ui/AIAssistantPanel/AIAssistant.types.ts

@@ -10,3 +10,5 @@ export interface AssistantSnippetProps {
   yAxis?: string
   name?: string
 }
+
+export type SqlSnippet = string | { label: string; content: string }

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

@@ -4,11 +4,12 @@ import { useRef } from 'react'
 
 import { Button, cn } from 'ui'
 import { AssistantChatForm } from './AssistantChatForm'
+import { type SqlSnippet } from './AIAssistant.types'
 import { codeSnippetPrompts, defaultPrompts } from './AIAssistant.prompts'
 
 interface AIOnboardingProps {
   onMessageSend: (message: string) => void
-  sqlSnippets?: string[]
+  sqlSnippets?: SqlSnippet[]
   onRemoveSnippet?: (index: number) => void
   suggestions?: {
     title?: string

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

@@ -5,7 +5,8 @@ import { ArrowUp, Loader2 } from 'lucide-react'
 import React, { ChangeEvent, memo, useRef } from 'react'
 import { Button, ExpandingTextArea } from 'ui'
 import { cn } from 'ui/src/lib/utils'
-import { SnippetRow } from './SnippetRow'
+import { SnippetRow, getSnippetContent } from './SnippetRow'
+import { type SqlSnippet } from './AIAssistant.types'
 
 export interface FormProps {
   /* The ref for the textarea, optional. Exposed for the CommandsPopover to attach events. */
@@ -28,8 +29,8 @@ export interface FormProps {
   onSubmit: (message: string) => void
   /* The placeholder of the textarea */
   placeholder?: string
-  /* SQL snippets to display above the form */
-  sqlSnippets?: string[]
+  /* SQL snippets to display above the form - can be strings or objects with label and content */
+  sqlSnippets?: SqlSnippet[]
   /* Function to handle removing a SQL snippet */
   onRemoveSnippet?: (index: number) => void
   /* Additional class name for the snippets container */
@@ -67,7 +68,7 @@ const AssistantChatFormComponent = React.forwardRef<HTMLFormElement, FormProps>(
       let finalMessage = value
       if (includeSnippetsInMessage && sqlSnippets && sqlSnippets.length > 0) {
         const sqlSnippetsString = sqlSnippets
-          .map((snippet: string) => '```sql\n' + snippet + '\n```')
+          .map((snippet: SqlSnippet) => '```sql\n' + getSnippetContent(snippet) + '\n```')
           .join('\n')
         finalMessage = [value, sqlSnippetsString].filter(Boolean).join('\n\n')
       }

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

@@ -121,7 +121,7 @@ export const Message = function Message({
                           <ReactMarkdown
                             key={`${id}-part-${index}`}
                             className={cn(
-                              'prose prose-sm prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground break-words [&>p:not(:last-child)]:!mb-2 [&>*>p:first-child]:!mt-0 [&>*>p:last-child]:!mb-0 [&>*>*>p:first-child]:!mt-0 [&>*>*>p:last-child]:!mb-0 [&>ol>li]:!pl-4',
+                              'prose prose-sm [&>div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground break-words [&>p:not(:last-child)]:!mb-2 [&>*>p:first-child]:!mt-0 [&>*>p:last-child]:!mb-0 [&>*>*>p:first-child]:!mt-0 [&>*>*>p:last-child]:!mb-0 [&>ol>li]:!pl-4',
                               isUser && 'text-foreground [&>p]:font-medium'
                             )}
                             remarkPlugins={[remarkGfm]}

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

@@ -2,9 +2,24 @@ import React from 'react'
 import { Button, CodeBlock } from 'ui'
 import { HoverCard_Shadcn_, HoverCardTrigger_Shadcn_, HoverCardContent_Shadcn_ } from 'ui'
 import { X } from 'lucide-react'
+import { type SqlSnippet } from './AIAssistant.types'
+
+export const getSnippetLabel = (snippet: SqlSnippet, index: number): string => {
+  if (typeof snippet === 'string') {
+    return `Snippet ${index + 1}`
+  }
+  return snippet.label
+}
+
+export const getSnippetContent = (snippet: SqlSnippet): string => {
+  if (typeof snippet === 'string') {
+    return snippet
+  }
+  return snippet.content
+}
 
 interface SnippetRowProps {
-  snippets: string[]
+  snippets: SqlSnippet[]
   onRemoveSnippet?: (index: number) => void
   className?: string
 }
@@ -17,15 +32,15 @@ export const SnippetRow: React.FC<SnippetRowProps> = ({
   if (!snippets || snippets.length === 0) return null
 
   return (
-    <div className={`w-full overflow-x-auto flex gap-2 ${className}`}>
+    <div className={`w-full overflow-x-auto flex ${className}`}>
       {snippets.map((snippet, idx) => (
         <HoverCard_Shadcn_ key={idx}>
           <HoverCardTrigger_Shadcn_ asChild>
             <div
               tabIndex={0}
-              className="border inline-flex gap-1 items-center shrink-0 py-1 pl-2 rounded-full pr-1 text-xs cursor-pointer"
+              className="border bg inline-flex gap-1 items-center shrink-0 py-1 pl-2 rounded-full pr-1 text-xs cursor-pointer"
             >
-              Snippet {idx + 1}
+              {getSnippetLabel(snippet, idx)}
               {onRemoveSnippet && (
                 <Button
                   size="tiny"
@@ -47,7 +62,7 @@ export const SnippetRow: React.FC<SnippetRowProps> = ({
               className="text-xs font-mono whitespace-pre-wrap break-words p-2 border-0"
               language="sql"
             >
-              {snippet}
+              {getSnippetContent(snippet)}
             </CodeBlock>
           </HoverCardContent_Shadcn_>
         </HoverCard_Shadcn_>

apps/studio/lib/helpers.test.ts

@@ -18,6 +18,7 @@ import {
   removeCommentsFromSql,
   removeJSONTrailingComma,
   snakeToCamel,
+  tablesToSQL,
   timeout,
   tryParseInt,
   tryParseJson,
@@ -350,3 +351,146 @@ describe('formatCurrency', () => {
     expect(result).toEqual(null)
   })
 })
+
+describe('tablesToSQL', () => {
+  it('should return warning message for empty array', () => {
+    const result = tablesToSQL([])
+
+    expect(result).toContain('-- WARNING: This schema is for context only')
+  })
+
+  it('should return empty string for non-array input', () => {
+    const result = tablesToSQL(null as any)
+
+    expect(result).toBe('')
+  })
+
+  it('should generate SQL for a simple table', () => {
+    const mockTables = [
+      {
+        name: 'users',
+        schema: 'public',
+        columns: [
+          {
+            name: 'id',
+            data_type: 'integer',
+            is_nullable: false,
+            is_identity: true,
+            default_value: null,
+            is_unique: false,
+            check: null,
+          },
+          {
+            name: 'name',
+            data_type: 'text',
+            is_nullable: false,
+            is_identity: false,
+            default_value: null,
+            is_unique: false,
+            check: null,
+          },
+        ],
+        primary_keys: [{ name: 'id' }],
+        relationships: [],
+      },
+    ] as any
+
+    const result = tablesToSQL(mockTables)
+
+    expect(result).toContain('-- WARNING: This schema is for context only')
+    expect(result).toContain('CREATE TABLE public.users (')
+    expect(result).toContain('id integer GENERATED ALWAYS AS IDENTITY NOT NULL')
+    expect(result).toContain('name text NOT NULL')
+    expect(result).toContain('CONSTRAINT users_pkey PRIMARY KEY (id)')
+  })
+
+  it('should handle tables with various column properties', () => {
+    const mockTables = [
+      {
+        name: 'products',
+        schema: 'public',
+        columns: [
+          {
+            name: 'id',
+            data_type: 'uuid',
+            is_nullable: false,
+            is_identity: false,
+            default_value: 'gen_random_uuid()',
+            is_unique: true,
+            check: null,
+          },
+          {
+            name: 'price',
+            data_type: 'numeric',
+            is_nullable: true,
+            is_identity: false,
+            default_value: '0.00',
+            is_unique: false,
+            check: 'price >= 0',
+          },
+        ],
+        primary_keys: [],
+        relationships: [],
+      },
+    ] as any
+
+    const result = tablesToSQL(mockTables)
+
+    expect(result).toContain('id uuid NOT NULL DEFAULT gen_random_uuid() UNIQUE')
+    expect(result).toContain('price numeric DEFAULT 0.00 CHECK (price >= 0)')
+  })
+
+  it('should handle foreign key relationships', () => {
+    const mockTables = [
+      {
+        name: 'orders',
+        schema: 'public',
+        columns: [
+          {
+            name: 'user_id',
+            data_type: 'integer',
+            is_nullable: false,
+            is_identity: false,
+            default_value: null,
+            is_unique: false,
+            check: null,
+          },
+        ],
+        primary_keys: [],
+        relationships: [
+          {
+            constraint_name: 'fk_orders_user_id',
+            source_table_name: 'orders',
+            source_column_name: 'user_id',
+            target_table_schema: 'public',
+            target_table_name: 'users',
+            target_column_name: 'id',
+          },
+        ],
+      },
+    ] as any
+
+    const result = tablesToSQL(mockTables)
+
+    expect(result).toContain(
+      'CONSTRAINT fk_orders_user_id FOREIGN KEY (user_id) REFERENCES public.users(id)'
+    )
+  })
+
+  it('should handle tables with no columns', () => {
+    const mockTables = [
+      {
+        name: 'empty_table',
+        schema: 'public',
+        columns: null,
+        primary_keys: [],
+        relationships: [],
+      },
+    ] as any
+
+    const result = tablesToSQL(mockTables)
+
+    expect(result).toContain('-- WARNING: This schema is for context only')
+    expect(result).not.toContain('CREATE TABLE')
+  })
+})

apps/studio/lib/helpers.ts

@@ -1,5 +1,6 @@
 export { default as passwordStrength } from './password-strength'
 export { default as uuidv4 } from './uuid'
+import type { TablesData } from '../data/tables/tables-query'
 
 export const tryParseJson = (jsonString: any) => {
   try {
@@ -172,6 +173,64 @@ export const detectOS = () => {
   }
 }
 
+/**
+ * Convert a list of tables to SQL
+ * @param t - The list of tables
+ * @returns The SQL string
+ */
+export function tablesToSQL(t: TablesData) {
+  if (!Array.isArray(t)) return ''
+  const warning =
+    '-- WARNING: This schema is for context only and is not meant to be run.\n-- Table order and constraints may not be valid for execution.\n\n'
+  const sql = t
+    .map((table) => {
+      if (!table || !Array.isArray((table as any).columns)) return ''
+
+      const columns = (table as { columns?: any[] }).columns ?? []
+      const columnLines = columns.map((c) => {
+        let line = `  ${c.name} ${c.data_type}`
+        if (c.is_identity) {
+          line += ' GENERATED ALWAYS AS IDENTITY'
+        }
+        if (c.is_nullable === false) {
+          line += ' NOT NULL'
+        }
+        if (c.default_value !== null && c.default_value !== undefined) {
+          line += ` DEFAULT ${c.default_value}`
+        }
+        if (c.is_unique) {
+          line += ' UNIQUE'
+        }
+        if (c.check) {
+          line += ` CHECK (${c.check})`
+        }
+        return line
+      })
+
+      const constraints: string[] = []
+
+      if (Array.isArray(table.primary_keys) && table.primary_keys.length > 0) {
+        const pkCols = table.primary_keys.map((pk: any) => pk.name).join(', ')
+        constraints.push(`  CONSTRAINT ${table.name}_pkey PRIMARY KEY (${pkCols})`)
+      }
+
+      if (Array.isArray(table.relationships)) {
+        table.relationships.forEach((rel: any) => {
+          if (rel && rel.source_table_name === table.name) {
+            constraints.push(
+              `  CONSTRAINT ${rel.constraint_name} FOREIGN KEY (${rel.source_column_name}) REFERENCES ${rel.target_table_schema}.${rel.target_table_name}(${rel.target_column_name})`
+            )
+          }
+        })
+      }
+
+      const allLines = [...columnLines, ...constraints]
+      return `CREATE TABLE ${table.schema}.${table.name} (\n${allLines.join(',\n')}\n);`
+    })
+    .join('\n')
+  return warning + sql
+}
+
 /**
  * Pluralize a word based on a count
  */

apps/studio/pages/api/ai/sql/generate-v4.ts

@@ -88,7 +88,9 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
   ])
 
   const selectedOrg = organizations.find((org) => org.slug === orgSlug)
-  const selectedProject = projects.find((project) => project.ref === projectRef)
+  const selectedProject = projects.find(
+    (project) => project.ref === projectRef || project.preview_branch_refs.includes(projectRef)
+  )
 
   // If the project is not in the organization specific by the org slug, return an error
   if (selectedProject?.organization_slug !== selectedOrg?.slug) {

apps/studio/pages/project/[ref]/merge.tsx

@@ -47,6 +47,7 @@ import { ButtonTooltip } from 'components/ui/ButtonTooltip'
 import WorkflowLogsCard from 'components/interfaces/BranchManagement/WorkflowLogsCard'
 import ProductEmptyState from 'components/to-be-cleaned/ProductEmptyState'
 import { useFlag } from 'hooks/ui/useFlag'
+import { ReviewWithAI } from 'components/interfaces/BranchManagement/ReviewWithAI'
 
 const MergePage: NextPageWithLayout = () => {
   const router = useRouter()
@@ -353,6 +354,13 @@ const MergePage: NextPageWithLayout = () => {
   // Update primary actions - remove push button if branch is out of date (it will be in the notice)
   const primaryActions = (
     <div className="flex items-end gap-2">
+      <ReviewWithAI
+        currentBranch={currentBranch}
+        mainBranch={mainBranch}
+        parentProjectRef={parentProjectRef}
+        diffContent={diffContent}
+        disabled={!currentBranch || !mainBranch || isCombinedDiffLoading}
+      />
       {!isReadyForReview ? (
         <Button
           type="primary"

apps/studio/state/ai-assistant-state.tsx

@@ -14,6 +14,8 @@ type SuggestionsType = {
 
 export type AssistantMessageType = MessageType & { results?: { [id: string]: any[] } }
 
+export type SqlSnippet = string | { label: string; content: string }
+
 type ChatSession = {
   id: string
   name: string
@@ -25,7 +27,7 @@ type ChatSession = {
 type AiAssistantData = {
   open: boolean
   initialInput: string
-  sqlSnippets?: string[]
+  sqlSnippets?: SqlSnippet[]
   suggestions?: SuggestionsType
   tables: { schema: string; name: string }[]
   chats: Record<string, ChatSession>
@@ -335,7 +337,7 @@ export const createAiAssistantState = (): AiAssistantState => {
       }
     },
 
-    setSqlSnippets: (snippets: string[]) => {
+    setSqlSnippets: (snippets: SqlSnippet[]) => {
       state.sqlSnippets = snippets
     },
 
@@ -406,7 +408,7 @@ export type AiAssistantState = AiAssistantData & {
   clearMessages: () => void
   saveMessage: (message: MessageType | MessageType[]) => void
   updateMessage: (args: { id: string; resultId?: string; results: any[] }) => void
-  setSqlSnippets: (snippets: string[]) => void
+  setSqlSnippets: (snippets: SqlSnippet[]) => void
   clearSqlSnippets: () => void
   getCachedSQLResults: (args: { messageId: string; snippetId?: string }) => any[] | undefined
   loadPersistedState: (persistedState: StoredAiAssistantState) => void