From 3229b554e8ba899a8356add32ef5de7e243fc519 Mon Sep 17 00:00:00 2001
From: Saxon Fletcher <saxonafletcher@gmail.com>
Date: Thu, 6 Mar 2025 08:53:01 +1000
Subject: [PATCH] refactor

---
 .../layouts/ProjectLayout/ProjectLayout.tsx   |  4 +---
 .../ui/AIAssistantPanel/AIAssistant.tsx       |  2 ++
 .../AIAssistantChatSelector.tsx               |  9 ++++-----
 apps/studio/hooks/useAssistant.ts             | 19 -------------------
 apps/studio/state/app-state.ts                | 10 +---------
 5 files changed, 8 insertions(+), 36 deletions(-)

diff --git a/apps/studio/components/layouts/ProjectLayout/ProjectLayout.tsx b/apps/studio/components/layouts/ProjectLayout/ProjectLayout.tsx
index af3ada7abfc..c205e59404c 100644
--- a/apps/studio/components/layouts/ProjectLayout/ProjectLayout.tsx
+++ b/apps/studio/components/layouts/ProjectLayout/ProjectLayout.tsx
@@ -228,9 +228,7 @@ const ProjectLayout = forwardRef<HTMLDivElement, PropsWithChildren<ProjectLayout
                       )}
                     >
                       {aiAssistantPanel.open && (
-                        <div className="w-full h-[100dvh] md:h-full max-h-[100dvh]">
-                          <AIAssistant className="w-full h-[100dvh] md:h-full max-h-[100dvh]" />
-                        </div>
+                        <AIAssistant className="w-full h-[100dvh] md:h-full max-h-[100dvh]" />
                       )}
                       {editorPanel.open && <EditorPanel />}
                     </ResizablePanel>
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
index 70c96211c8f..9c76b3713d6 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
@@ -85,6 +85,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     saveMessages,
     clearMessages,
     closeAssistant,
+    clearSqlSnippets,
     newChat,
     clearSuggestions,
     setSqlSnippets,
@@ -176,6 +177,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   const sendMessageToAssistant = async (content: string) => {
     const payload = { role: 'user', createdAt: new Date(), content } as MessageType
     const headerData = await constructHeaders()
+    clearSqlSnippets()
     append(payload, {
       headers: { Authorization: headerData.get('Authorization') ?? '' },
     })
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx
index c2fab4b3f46..ad60060f628 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx
@@ -1,5 +1,5 @@
 import { useState } from 'react'
-import { Check, ChevronsUpDown, Edit, MessageSquare, Plus, Trash, X } from 'lucide-react'
+import { Check, Edit, MessageSquare, Trash, X } from 'lucide-react'
 import { useAssistant } from 'hooks/useAssistant'
 import { useRouter } from 'next/router'
 import {
@@ -11,7 +11,6 @@ import {
   CommandInput_Shadcn_,
   CommandItem_Shadcn_,
   CommandList_Shadcn_,
-  Input,
   Input_Shadcn_,
   Popover_Shadcn_,
   PopoverContent_Shadcn_,
@@ -25,7 +24,7 @@ interface AIAssistantChatSelectorProps {
 export const AIAssistantChatSelector = ({ className }: AIAssistantChatSelectorProps) => {
   const router = useRouter()
   const projectRef = typeof router.query.ref === 'string' ? router.query.ref : undefined
-  const { chats, activeChatId, selectChat, deleteChat, renameChat, newChat } = useAssistant({
+  const { chats, activeChatId, selectChat, deleteChat, renameChat } = useAssistant({
     projectRef,
   })
 
@@ -105,7 +104,7 @@ export const AIAssistantChatSelector = ({ className }: AIAssistantChatSelectorPr
                   key={id}
                   value={id}
                   onSelect={() => handleSelectChat(id)}
-                  className="flex items-center justify-between gap-2 py-1 w-full overflow-hidden"
+                  className="flex items-center justify-between gap-2 py-1 w-full overflow-hidden group"
                   keywords={[chat.name]}
                 >
                   <div className="flex items-center w-full flex-1 min-w-0">
@@ -161,7 +160,7 @@ export const AIAssistantChatSelector = ({ className }: AIAssistantChatSelectorPr
                     )}
                   </div>
                   {editingChatId !== id && (
-                    <div className="flex items-center gap-0 shrink-0">
+                    <div className="flex items-center gap-0 shrink-0 opacity-0 group-hover:opacity-100 transition-opacity">
                       <Button
                         type="text"
                         size="tiny"
diff --git a/apps/studio/hooks/useAssistant.ts b/apps/studio/hooks/useAssistant.ts
index b070d1b0156..79fa074fb5f 100644
--- a/apps/studio/hooks/useAssistant.ts
+++ b/apps/studio/hooks/useAssistant.ts
@@ -3,23 +3,6 @@ import { useCallback, useEffect, useMemo } from 'react'
 import { useAppStateSnapshot } from 'state/app-state'
 import { ChatSession } from 'state/app-state'
 
-// Define a simpler type for messages to avoid deep instantiation errors
-// Make it compatible with the Message type from ai/react
-type SimpleMessage = {
-  id: string
-  content: string
-  role: 'function' | 'data' | 'system' | 'user' | 'assistant' | 'tool' | string
-  [key: string]: any
-}
-
-// Helper function to convert SimpleMessage[] to MessageType[]
-const convertToMessageType = (messages: SimpleMessage[]): MessageType[] => {
-  return messages.map((msg) => ({
-    ...msg,
-    role: msg.role as 'function' | 'data' | 'system' | 'user' | 'assistant' | 'tool',
-  })) as MessageType[]
-}
-
 export interface UseAssistantOptions {
   projectRef?: string
 }
@@ -215,7 +198,6 @@ export function useAssistant(options?: UseAssistantOptions) {
 
   const handleSaveMessage = useCallback(
     (message: MessageType) => {
-      clearSqlSnippets()
       if (projectRef && currentChatBelongsToProject && activeChatId) {
         // Get the existing messages for the active chat
         const existingMessages = projectChatsRecord[activeChatId]?.messages || []
@@ -253,7 +235,6 @@ export function useAssistant(options?: UseAssistantOptions) {
 
   const handleSaveMessages = useCallback(
     (messages: MessageType[]) => {
-      clearSqlSnippets()
       if (projectRef && currentChatBelongsToProject && activeChatId) {
         // Get the current messages to compare
         const currentMessages = projectChatsRecord[activeChatId]?.messages || []
diff --git a/apps/studio/state/app-state.ts b/apps/studio/state/app-state.ts
index 4596ef91e9e..c0df5ce2525 100644
--- a/apps/studio/state/app-state.ts
+++ b/apps/studio/state/app-state.ts
@@ -122,9 +122,6 @@ const getInitialState = () => {
   const urlParams = new URLSearchParams(window.location.search)
   const aiAssistantPanelOpenParam = urlParams.get('aiAssistantPanelOpen')
 
-  // Get project ref from URL if available
-  const projectRef = urlParams.get('ref') || ''
-
   let parsedAiAssistant = INITIAL_AI_ASSISTANT
   let parsedEditorPanel = INITIAL_EDITOR_PANEL
 
@@ -283,9 +280,6 @@ export const appState = proxy({
   setMobileMenuOpen: (value: boolean) => {
     appState.mobileMenuOpen = value
   },
-
-  // The assistant-related functions have been moved to the useAssistant hook
-  // Only keeping the setAiAssistantPanel function for state updates
 })
 
 // Set up localStorage subscriptions
@@ -294,23 +288,21 @@ if (typeof window !== 'undefined') {
     // Save AI assistant state with limited message history
     const aiAssistantState = {
       ...appState.aiAssistantPanel,
-      // Process chats to limit message history in each chat
       chats: appState.aiAssistantPanel.chats
         ? Object.entries(appState.aiAssistantPanel.chats).reduce((acc, [chatId, chat]) => {
             return {
               ...acc,
               [chatId]: {
                 ...chat,
+                messages: chat.messages?.slice(-20) || [], // Only keep last 20 messages
               },
             }
           }, {})
         : {},
     }
 
-    // Save to global storage
     localStorage.setItem(LOCAL_STORAGE_KEYS.AI_ASSISTANT_STATE, JSON.stringify(aiAssistantState))
 
-    // Save editor panel state
     localStorage.setItem(
       LOCAL_STORAGE_KEYS.EDITOR_PANEL_STATE,
       JSON.stringify(appState.editorPanel)