refactor to use onfinish

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

@@ -82,7 +82,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   const { aiAssistantPanel } = useAppStateSnapshot()
   const {
     messages: assistantMessages,
-    saveMessages,
+    saveMessage,
     clearMessages,
     closeAssistant,
     clearSqlSnippets,
@@ -96,6 +96,9 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   const inputRef = useRef<HTMLTextAreaElement>(null)
   const { ref: scrollContainerRef, isSticky, scrollToEnd } = useAutoScroll()
 
+  // Add a ref to store the last user message
+  const lastUserMessageRef = useRef<MessageType | null>(null)
+
   const [value, setValue] = useState<string>(initialInput)
   const [isConfirmOptInModalOpen, setIsConfirmOptInModalOpen] = useState(false)
 
@@ -130,6 +133,21 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     toast.error(errorMessage)
   }, [])
 
+  // Handle completion of the assistant's response
+  const handleChatFinish = useCallback(
+    (message: MessageType) => {
+      // If we have a user message stored in the ref, save both messages
+      if (lastUserMessageRef.current) {
+        saveMessage([lastUserMessageRef.current, message])
+        lastUserMessageRef.current = null
+      } else {
+        // Otherwise just save the assistant message
+        saveMessage(message)
+      }
+    },
+    [saveMessage]
+  )
+
   const {
     messages: chatMessages,
     isLoading: isChatLoading,
@@ -148,13 +166,9 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
       table: currentTable?.name,
     },
     onError: handleError,
+    onFinish: handleChatFinish,
   })
 
-  // Watch messages and save back to global state when changes
-  useEffect(() => {
-    saveMessages(chatMessages)
-  }, [chatMessages, saveMessages])
-
   const canUpdateOrganization = useCheckPermissions(PermissionAction.UPDATE, 'organizations')
   const { mutate: updateOrganization, isLoading: isUpdating } = useOrganizationUpdateMutation()
 
@@ -178,6 +192,10 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     const payload = { role: 'user', createdAt: new Date(), content } as MessageType
     const headerData = await constructHeaders()
     clearSqlSnippets()
+
+    // Store the user message in the ref before appending
+    lastUserMessageRef.current = payload
+
     append(payload, {
       headers: { Authorization: headerData.get('Authorization') ?? '' },
     })
@@ -224,6 +242,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   const handleClearMessages = () => {
     clearMessages()
     setMessages([])
+    lastUserMessageRef.current = null
   }
 
   // Update scroll behavior for new messages
@@ -235,7 +254,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     if (isSticky) {
       setTimeout(scrollToEnd, 0)
     }
-  }, [isChatLoading, isSticky, scrollToEnd, chatMessages])
+  }, [isChatLoading, isSticky, scrollToEnd])
 
   useEffect(() => {
     setValue(initialInput)

apps/studio/hooks/useAssistant.ts

@@ -113,12 +113,11 @@ export function useAssistant(options?: UseAssistantOptions) {
 
   const handleSelectChat = useCallback(
     (id: string) => {
-      if (projectRef && chats && id in chats) {
-        const chat = projectChatsRecord[id]
-        if (chat) {
-          // Check if this is already the active chat
-          if (id === activeChatId) {
-            // If it's already the active chat, no need to update state
+      if (id !== activeChatId) {
+        if (projectRef && projectChatsRecord) {
+          // Check if the selected chat exists and belongs to the current project
+          const targetChat = projectChatsRecord[id]
+          if (!targetChat) {
             return
           }
 
@@ -128,7 +127,7 @@ export function useAssistant(options?: UseAssistantOptions) {
         }
       }
     },
-    [activeChatId, chats, projectChatsRecord, projectRef, setAiAssistantPanel]
+    [activeChatId, projectRef, setAiAssistantPanel]
   )
 
   const handleDeleteChat = useCallback(
@@ -197,61 +196,32 @@ export function useAssistant(options?: UseAssistantOptions) {
   }, [activeChatId, chats, currentChatBelongsToProject, projectChatsRecord, setAiAssistantPanel])
 
   const handleSaveMessage = useCallback(
-    (message: MessageType) => {
+    (message: MessageType | MessageType[]) => {
       if (projectRef && currentChatBelongsToProject && activeChatId) {
         // Get the existing messages for the active chat
-        const existingMessages = projectChatsRecord[activeChatId]?.messages || []
-        const messageExists = existingMessages.some(
-          (msg: any) =>
-            msg.id === message.id || (msg.content === message.content && msg.role === message.role)
-        )
+        const existingMessages: readonly MessageType[] =
+          projectChatsRecord[activeChatId]?.messages || []
 
-        if (!messageExists) {
-          const newMessages = [...existingMessages, message]
+        // Convert single message to array for consistent handling
+        const messagesToAdd = Array.isArray(message)
+          ? message.filter((msg) => !existingMessages.some((existing) => existing.id === msg.id))
+          : !existingMessages.some((existing) => existing.id === message.id)
+            ? [message]
+            : []
 
-          // Update chat messages
-          setAiAssistantPanel({
+        if (messagesToAdd.length > 0) {
+          const chatUpdate = {
             chats: {
               ...chats,
               [activeChatId]: {
                 ...projectChatsRecord[activeChatId],
-                messages: newMessages,
+                messages: [...existingMessages, ...messagesToAdd],
                 updatedAt: new Date(),
               },
-            } as any,
-          })
-        }
-      }
-    },
-    [
-      chats,
-      activeChatId,
-      currentChatBelongsToProject,
-      projectRef,
-      projectChatsRecord,
-      setAiAssistantPanel,
-    ]
-  )
-
-  const handleSaveMessages = useCallback(
-    (messages: MessageType[]) => {
-      if (projectRef && currentChatBelongsToProject && activeChatId) {
-        // Get the current messages to compare
-        const currentMessages = projectChatsRecord[activeChatId]?.messages || []
-
-        // Check if the messages array is actually different before updating
-        // This prevents circular updates
-        if (JSON.stringify(currentMessages) !== JSON.stringify(messages)) {
-          // Update chat messages
+            },
+          }
           setAiAssistantPanel({
-            chats: {
-              ...chats,
-              [activeChatId]: {
-                ...projectChatsRecord[activeChatId],
-                messages,
-                updatedAt: new Date(),
-              },
-            } as any,
+            chats: chatUpdate.chats as Record<string, ChatSession>,
           })
         }
       }
@@ -336,7 +306,6 @@ export function useAssistant(options?: UseAssistantOptions) {
     renameChat: handleRenameChat,
     clearMessages: handleClearMessages,
     saveMessage: handleSaveMessage,
-    saveMessages: handleSaveMessages,
     openAssistant,
     closeAssistant,
     clearSqlSnippets,

apps/studio/state/app-state.ts

@@ -27,7 +27,7 @@ export type ChatSession = {
   id: string
   name: string
   projectRef: string
-  messages: MessageType[]
+  messages: readonly MessageType[]
   createdAt: Date
   updatedAt: Date
 }