fixes

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

@@ -13,6 +13,8 @@ import { SIDEBAR_KEYS } from 'components/layouts/ProjectLayout/LayoutSidebar/Lay
 import { useCheckOpenAIKeyQuery } from 'data/ai/check-api-key-query'
 import { useRateMessageMutation } from 'data/ai/rate-message-mutation'
 import { useChatSessionCreateMutation } from 'data/chat-sessions/chat-session-create-mutation'
+import { useChatSessionMessageDeleteMutation } from 'data/chat-sessions/chat-session-message-delete-mutation'
+import { useChatSessionMessagesCreateMutation } from 'data/chat-sessions/chat-session-messages-create-mutation'
 import { useChatSessionMessagesQuery } from 'data/chat-sessions/chat-session-messages-query'
 import { useChatSessionsQuery } from 'data/chat-sessions/chat-sessions-query'
 import { useTablesQuery } from 'data/tables/tables-query'
@@ -67,13 +69,15 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   const { activeSidebar, closeSidebar } = useSidebarManagerSnapshot()
 
   // Avoid hammering the API when project ref is unavailable or if previous calls failed
-  const shouldFetchSessions = IS_PLATFORM && !!project?.ref && !disablePrompts
   const [sessionFetchError, setSessionFetchError] = useState<string | null>(null)
+  const shouldFetchSessions =
+    IS_PLATFORM && !!project?.ref && !disablePrompts && sessionFetchError === null
 
   // Load sessions from server
   const {
     data: sessions = [],
-    isLoading: isLoadingSessions,
+    isPending: isPendingSessions,
+    isSuccess: isSuccessSessions,
     error: chatSessionsError,
   } = useChatSessionsQuery(
     { projectRef: project?.ref },
@@ -87,12 +91,24 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   )
 
   // Load messages for active chat
-  const { data: initialMessages = [] } = useChatSessionMessagesQuery(
+  const {
+    data: initialMessages = [],
+    isSuccess: isSuccessInitialMessages,
+  } = useChatSessionMessagesQuery<MessageType[]>(
     { id: snap.activeChatId, projectRef: project?.ref },
-    { enabled: shouldFetchSessions && !!snap.activeChatId && !chatSessionsError }
+    {
+      enabled:
+        shouldFetchSessions &&
+        isSuccessSessions &&
+        !!snap.activeChatId &&
+        !chatSessionsError &&
+        typeof project?.ref !== 'undefined',
+    }
   )
 
   const createSession = useChatSessionCreateMutation()
+  const saveMessages = useChatSessionMessagesCreateMutation()
+  const deleteMessageMutation = useChatSessionMessageDeleteMutation()
 
   const [updatedOptInSinceMCP] = useLocalStorageQuery(
     LOCAL_STORAGE_KEYS.AI_ASSISTANT_MCP_OPT_IN,
@@ -113,6 +129,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
   }, [isPaidPlan, snap.model])
 
   const inputRef = useRef<HTMLTextAreaElement>(null)
+  const hasHydratedChatMessagesRef = useRef<string | null>(null)
 
   const { aiOptInLevel, isHipaaProjectDisallowed } = useOrgAiOptInLevel()
   const showMetadataWarning =
@@ -120,39 +137,42 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
     !!selectedOrganization &&
     (aiOptInLevel === 'disabled' || aiOptInLevel === 'schema')
 
-  // // Initialize active chat if none selected
-  // useEffect(() => {
-  //   if (!shouldFetchSessions || isLoadingSessions || chatSessionsError) return
-  //   if (!snap.activeChatId && sessions.length > 0) {
-  //     // Select the most recently updated chat
-  //     const mostRecentSession = sessions.sort(
-  //       (a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()
-  //     )[0]
-  //     state.setActiveChatId(mostRecentSession.id)
-  //   } else if (!snap.activeChatId && sessions.length === 0 && project?.ref) {
-  //     // Create a new chat if none exist
-  //     createSession.mutate(
-  //       { projectRef: project.ref },
-  //       {
-  //         onSuccess: (newSession) => {
-  //           state.setActiveChatId(newSession.id)
-  //         },
-  //         onError: (err) => {
-  //           setSessionFetchError(err.message ?? 'Failed to create chat session')
-  //         },
-  //       }
-  //     )
-  //   }
-  // }, [
-  //   shouldFetchSessions,
-  //   isLoadingSessions,
-  //   chatSessionsError,
-  //   sessions,
-  //   snap.activeChatId,
-  //   project?.ref,
-  //   state,
-  //   createSession,
-  // ])
+  // Initialize active chat if none selected (global state lives in ai-assistant-state.tsx)
+  useEffect(() => {
+    if (!shouldFetchSessions || isPendingSessions || chatSessionsError) return
+    if (snap.activeChatId) return
+
+    if (sessions.length > 0) {
+      const mostRecentSession = [...sessions].sort(
+        (a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()
+      )[0]
+      state.setActiveChatId(mostRecentSession.id)
+      return
+    }
+
+    if (sessions.length === 0 && project?.ref) {
+      createSession.mutate(
+        { projectRef: project.ref },
+        {
+          onSuccess: (newSession) => {
+            state.setActiveChatId(newSession.id)
+          },
+          onError: (err) => {
+            setSessionFetchError(err.message ?? 'Failed to create chat session')
+          },
+        }
+      )
+    }
+  }, [
+    shouldFetchSessions,
+    isPendingSessions,
+    chatSessionsError,
+    sessions,
+    snap.activeChatId,
+    project?.ref,
+    state,
+    createSession,
+  ])
 
   // Keep latest selected organization to avoid stale values in useChat transport
   const selectedOrganizationRef = useRef(selectedOrganization)
@@ -218,6 +238,33 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
 
   const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming'
 
+  // Hydrate the current chat with server messages once per chat id.
+  useEffect(() => {
+    if (!snap.activeChatId) return
+    if (!shouldFetchSessions) return
+    if (!isSuccessInitialMessages) return
+    if (hasHydratedChatMessagesRef.current === snap.activeChatId) return
+    if (isChatLoading) return
+
+    if (initialMessages.length > 0) {
+      const mergedMessages = [...initialMessages, ...chatMessages]
+      const uniqueMessages = mergedMessages.filter(
+        (msg, idx, all) => idx === all.findIndex((m) => m.id === msg.id)
+      )
+      setMessages(uniqueMessages)
+    }
+
+    hasHydratedChatMessagesRef.current = snap.activeChatId
+  }, [
+    snap.activeChatId,
+    shouldFetchSessions,
+    isSuccessInitialMessages,
+    isChatLoading,
+    chatMessages,
+    initialMessages,
+    setMessages,
+  ])
+
   const deleteMessageFromHere = useCallback(
     (messageId: string) => {
       // Find the message index in current chatMessages
@@ -226,12 +273,24 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
 
       if (isChatLoading) stop()
 
-      snap.deleteMessagesAfter(messageId, { includeSelf: true })
-
+      // Get messages to delete (from messageId onwards)
+      const messagesToDelete = chatMessages.slice(messageIndex)
       const updatedMessages = chatMessages.slice(0, messageIndex)
+
       setMessages(updatedMessages)
+
+      // Delete messages from server
+      if (project?.ref && snap.activeChatId) {
+        for (const msg of messagesToDelete) {
+          deleteMessageMutation.mutate({
+            projectRef: project.ref,
+            chatId: snap.activeChatId,
+            messageId: msg.id,
+          })
+        }
+      }
     },
-    [snap, setMessages, chatMessages, isChatLoading, stop]
+    [chatMessages, isChatLoading, project?.ref, deleteMessageMutation, setMessages, snap.activeChatId, stop]
   )
 
   const editMessage = useCallback(
@@ -359,10 +418,8 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
 
   const hasMessages = chatMessages.length > 0
 
-  const sendMessageToAssistant = (finalContent: string) => {
+  const sendMessageToAssistant = async (finalContent: string) => {
     if (editingMessageId) {
-      // Handling when the user is in edit mode
-      // delete the message(s) from the chat just like the delete button
       setIsResubmitting(true)
       deleteMessageFromHere(editingMessageId)
       setEditingMessageId(null)
@@ -656,15 +713,20 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
             }
             value={value}
             onValueChange={(e) => setValue(e.target.value)}
-            onSubmit={(finalMessage) => {
-              sendMessageToAssistant(finalMessage)
-            }}
+            onSubmit={(finalMessage) => void sendMessageToAssistant(finalMessage)}
             onStop={() => {
               stop()
-              // to save partial responses from the AI
+              // Save partial assistant response to the database
               const lastMessage = chatMessages[chatMessages.length - 1]
               if (lastMessage && lastMessage.role === 'assistant') {
-                state.updateMessage(lastMessage)
+                if (project?.ref && snap.activeChatId) {
+                  // Save just the partial assistant message (user message already saved by server)
+                  saveMessages.mutate({
+                    projectRef: project.ref,
+                    id: snap.activeChatId,
+                    messages: [lastMessage] as any,
+                  })
+                }
               }
             }}
             sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}

apps/studio/data/chat-sessions/chat-session-message-delete-mutation.ts

@@ -0,0 +1,47 @@
+import { useMutation, useQueryClient } from '@tanstack/react-query'
+import { toast } from 'sonner'
+
+import type { ResponseError, UseCustomMutationOptions } from 'types'
+import { chatSessionKeys } from './keys'
+import {
+  type ChatSessionMessageDeleteData,
+  deleteChatSessionMessage,
+  type ChatSessionMessageDeleteVariables,
+} from './chat-session-message-delete'
+
+export const useChatSessionMessageDeleteMutation = ({
+  onSuccess,
+  onError,
+  ...options
+}: Omit<
+  UseCustomMutationOptions<
+    ChatSessionMessageDeleteData,
+    ResponseError,
+    ChatSessionMessageDeleteVariables
+  >,
+  'mutationFn'
+> = {}) => {
+  const queryClient = useQueryClient()
+
+  return useMutation<
+    ChatSessionMessageDeleteData,
+    ResponseError,
+    ChatSessionMessageDeleteVariables
+  >({
+    mutationFn: (vars) => deleteChatSessionMessage(vars),
+    async onSuccess(data, variables, context) {
+      await queryClient.invalidateQueries({
+        queryKey: chatSessionKeys.messages(variables.chatId),
+      })
+      await onSuccess?.(data, variables, context)
+    },
+    async onError(data, variables, context) {
+      if (onError === undefined) {
+        toast.error(`Failed to delete message: ${data.message}`)
+      } else {
+        onError(data, variables, context)
+      }
+    },
+    ...options,
+  })
+}

apps/studio/data/chat-sessions/chat-session-message-delete.ts

@@ -0,0 +1,34 @@
+import { del, handleError } from 'data/fetchers'
+import type { ResponseError } from 'types'
+
+export type ChatSessionMessageDeleteVariables = {
+  projectRef: string
+  chatId: string
+  messageId: string
+}
+
+export async function deleteChatSessionMessage(
+  { projectRef, chatId, messageId }: ChatSessionMessageDeleteVariables,
+  headersInit?: HeadersInit
+) {
+  if (!projectRef) throw new Error('projectRef is required')
+  if (!chatId) throw new Error('chatId is required')
+  if (!messageId) throw new Error('messageId is required')
+
+  const headers = new Headers(headersInit)
+
+  // Note: This endpoint path will be added to the API types when backend is updated
+  const { data, error } = await (del as any)(
+    `/v1/projects/{ref}/chat-sessions/{id}/messages/{messageId}`,
+    {
+      params: { path: { ref: projectRef, id: chatId, messageId } },
+      headers,
+    }
+  )
+
+  if (error) handleError(error)
+  return data
+}
+
+export type ChatSessionMessageDeleteData = Awaited<ReturnType<typeof deleteChatSessionMessage>>
+export type ChatSessionMessageDeleteError = ResponseError

apps/studio/data/chat-sessions/chat-session-messages-create-mutation.ts

@@ -0,0 +1,56 @@
+import { useMutation, useQueryClient } from '@tanstack/react-query'
+import { toast } from 'sonner'
+
+import type { ResponseError, UseCustomMutationOptions } from 'types'
+import { chatSessionKeys } from './keys'
+import {
+  type ChatSessionMessagesCreateData,
+  createChatSessionMessages,
+  type ChatSessionMessagesCreateVariables,
+} from './chat-session-messages-create'
+
+export const useChatSessionMessagesCreateMutation = ({
+  onSuccess,
+  onError,
+  ...options
+}: Omit<
+  UseCustomMutationOptions<
+    ChatSessionMessagesCreateData,
+    ResponseError,
+    ChatSessionMessagesCreateVariables
+  >,
+  'mutationFn'
+> = {}) => {
+  const queryClient = useQueryClient()
+
+  return useMutation<
+    ChatSessionMessagesCreateData,
+    ResponseError,
+    ChatSessionMessagesCreateVariables
+  >({
+    mutationFn: (vars) => createChatSessionMessages(vars),
+    async onSuccess(data, variables, context) {
+      await Promise.all([
+        queryClient.invalidateQueries({
+          queryKey: chatSessionKeys.messages(variables.id),
+        }),
+        queryClient.invalidateQueries({
+          queryKey: chatSessionKeys.detail(variables.id),
+        }),
+        queryClient.invalidateQueries({
+          queryKey: chatSessionKeys.list(variables.projectRef),
+        }),
+      ])
+      await onSuccess?.(data, variables, context)
+    },
+    async onError(data, variables, context) {
+      if (onError === undefined) {
+        toast.error(`Failed to save chat messages: ${data.message}`)
+      } else {
+        onError(data, variables, context)
+      }
+    },
+    ...options,
+  })
+}
+

apps/studio/data/chat-sessions/chat-session-messages-create.ts

@@ -0,0 +1,31 @@
+import type { UIMessage } from 'ai'
+import { handleError, post } from 'data/fetchers'
+import type { ResponseError } from 'types'
+
+export type ChatSessionMessagesCreateVariables = {
+  projectRef: string
+  id: string
+  messages: UIMessage[]
+}
+
+export async function createChatSessionMessages(
+  { projectRef, id, messages }: ChatSessionMessagesCreateVariables,
+  headersInit?: HeadersInit
+) {
+  if (!projectRef) throw new Error('projectRef is required')
+  if (!id) throw new Error('id is required')
+
+  const headers = new Headers(headersInit)
+
+  const { data, error } = await post(`/v1/projects/{ref}/chat-sessions/{id}/messages`, {
+    params: { path: { ref: projectRef, id } },
+    body: { messages: messages as any },
+    headers,
+  })
+
+  if (error) handleError(error)
+  return data
+}
+
+export type ChatSessionMessagesCreateData = Awaited<ReturnType<typeof createChatSessionMessages>>
+export type ChatSessionMessagesCreateError = ResponseError

apps/studio/data/chat-sessions/chat-session-messages-query.ts

@@ -1,4 +1,5 @@
 import { useQuery } from '@tanstack/react-query'
+import type { UIMessage } from 'ai'
 import { components } from 'api-types'
 import { get, handleError } from 'data/fetchers'
 import { IS_PLATFORM } from 'lib/constants'
@@ -14,18 +15,33 @@ export type ChatMessageResponse = components['schemas']['ChatMessage']
 
 export async function getChatSessionMessages(
   { id, projectRef }: ChatSessionMessagesVariables,
-  signal?: AbortSignal
+  signal?: AbortSignal,
+  headersInit?: HeadersInit
 ) {
   if (!id) throw new Error('id is required')
   if (!projectRef) throw new Error('projectRef is required')
 
+  const headers = new Headers(headersInit)
+
   const { data, error } = await get(`/v1/projects/{ref}/chat-sessions/{id}/messages`, {
     params: { path: { ref: projectRef, id } },
+    headers,
     signal,
   })
 
   if (error) handleError(error)
-  return data
+
+  // API returns { data: ChatMessage[] }
+  if (data && typeof data === 'object' && 'data' in data && Array.isArray((data as any).data)) {
+    return (data as any).data as UIMessage[]
+  }
+
+  // Fallback for direct array response
+  if (Array.isArray(data)) {
+    return data as UIMessage[]
+  }
+
+  throw new Error('Unexpected chat session messages response')
 }
 
 export type ChatSessionMessagesData = Awaited<ReturnType<typeof getChatSessionMessages>>

apps/studio/data/chat-sessions/chat-session-update-mutation.ts

@@ -10,13 +10,19 @@ export type ChatSessionUpdateVariables = {
   name?: string
 }
 
-export async function updateChatSession({ id, projectRef, name }: ChatSessionUpdateVariables) {
+export async function updateChatSession(
+  { id, projectRef, name }: ChatSessionUpdateVariables,
+  headersInit?: HeadersInit
+) {
   if (!id) throw new Error('id is required')
   if (!projectRef) throw new Error('projectRef is required')
 
+  const headers = headersInit ? new Headers(headersInit) : undefined
+
   const { data, error } = await patch(`/v1/projects/{ref}/chat-sessions/{id}`, {
     params: { path: { ref: projectRef, id } },
     body: { name },
+    ...(headers && { headers }),
   })
 
   if (error) handleError(error)

apps/studio/data/chat-sessions/chat-sessions-query.ts

@@ -20,8 +20,14 @@ export async function getChatSessions({ projectRef }: ChatSessionsVariables, sig
   })
 
   if (error) handleError(error)
-  console.log('data', data)
-  return data
+
+  if (Array.isArray(data)) return data as ChatSessionResponse[]
+
+  if (data && typeof data === 'object' && 'data' in data && Array.isArray((data as any).data)) {
+    return (data as any).data as ChatSessionResponse[]
+  }
+
+  throw new Error('Unexpected chat sessions response')
 }
 
 export type ChatSessionsData = Awaited<ReturnType<typeof getChatSessions>>

apps/studio/lib/ai/clean-message.ts

@@ -0,0 +1,39 @@
+import type { UIMessage } from 'ai'
+
+/**
+ * Cleans a UIMessage by removing streaming-only parts before persisting to database.
+ *
+ * Only removes:
+ * - step-start/step-finish parts (streaming boundary markers only useful during streaming)
+ * - tool parts in input-streaming state (incomplete - still being typed)
+ *
+ * Keeps all other parts including tool parts with their full context.
+ */
+export function cleanMessage(msg: UIMessage): UIMessage {
+  const cleanedParts = msg.parts.filter((part) => {
+    // Drop step-start/step-finish (streaming boundaries only)
+    if (part.type === 'step-start' || (part as { type: string }).type === 'step-finish') return false
+
+    // For tool parts, only drop if still in input-streaming state
+    if (part.type.startsWith('tool-') || part.type === 'dynamic-tool') {
+      if ((part as any).state === 'input-streaming') return false
+    }
+
+    // Keep everything else (text, reasoning, complete tool parts, etc.)
+    return true
+  })
+
+  return {
+    id: msg.id,
+    role: msg.role,
+    parts: cleanedParts,
+    ...(msg.metadata != null && typeof msg.metadata === 'object' ? { metadata: msg.metadata } : {}),
+  }
+}
+
+/**
+ * Cleans an array of UIMessages.
+ */
+export function cleanMessages(messages: UIMessage[]): UIMessage[] {
+  return messages.map(cleanMessage)
+}

apps/studio/lib/ai/normalize-ui-messages.test.ts

@@ -0,0 +1,77 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import { normalizeUiMessages } from './normalize-ui-messages'
+
+describe('normalizeUiMessages', () => {
+  it('drops streaming-only/invalid assistant parts for model input', () => {
+    const messages = [
+      {
+        id: 'user-1',
+        role: 'user',
+        parts: [{ type: 'text', text: 'Hello' }],
+      },
+      {
+        id: 'assistant-1',
+        role: 'assistant',
+        parts: [
+          { type: 'step-start' },
+          { type: 'reasoning' },
+          { type: 'tool-rename_chat' },
+          { type: 'text', text: 'Hi!' },
+        ],
+      },
+      { id: 'assistant-empty', role: 'assistant', parts: [] },
+    ]
+
+    const normalized = normalizeUiMessages(messages, {
+      dropReasoningParts: true,
+      dropToolStates: ['input-streaming', 'input-available', 'output-error'],
+      dropNonTextParts: true,
+    })
+
+    expect(normalized).toEqual([
+      { id: 'user-1', role: 'user', parts: [{ type: 'text', text: 'Hello' }] },
+      { id: 'assistant-1', role: 'assistant', parts: [{ type: 'text', text: 'Hi!' }] },
+    ])
+  })
+
+  it('keeps valid tool parts and runs the optional sanitizer', () => {
+    const sanitizer = vi.fn((part) => part)
+    const messages = [
+      {
+        id: 'assistant-1',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'tool-execute_sql',
+            toolCallId: 'call-1',
+            state: 'output-available',
+            input: { sql: 'select 1' },
+            output: [{ ok: true }],
+          },
+          { type: 'tool-execute_sql', state: 'output-available', output: [{ ok: true }] }, // missing toolCallId
+        ],
+      },
+    ]
+
+    const normalized = normalizeUiMessages(messages, { sanitizePart: sanitizer })
+
+    expect(normalized[0].parts).toHaveLength(1)
+    expect((normalized[0].parts[0] as any).toolCallId).toBe('call-1')
+    expect(sanitizer).toHaveBeenCalledTimes(1)
+  })
+
+  it('keeps well-formed reasoning parts when configured', () => {
+    const messages = [
+      {
+        id: 'assistant-1',
+        role: 'assistant',
+        parts: [{ type: 'reasoning', state: 'done', text: 'Thinking...' }, { type: 'text', text: 'Ok' }],
+      },
+    ]
+
+    const normalized = normalizeUiMessages(messages, { dropReasoningParts: false })
+
+    expect(normalized[0].parts[0]).toMatchObject({ type: 'reasoning', state: 'done', text: 'Thinking...' })
+  })
+})

apps/studio/lib/ai/normalize-ui-messages.ts

@@ -0,0 +1,165 @@
+import type { UIMessage } from 'ai'
+
+type UnknownRecord = Record<string, unknown>
+
+function isRecord(value: unknown): value is UnknownRecord {
+  return value != null && typeof value === 'object' && !Array.isArray(value)
+}
+
+export type NormalizeUiMessagesOptions = {
+  dropReasoningParts?: boolean
+  dropStepParts?: boolean
+  dropEmptyMessages?: boolean
+  dropToolStates?: string[]
+  dropNonTextParts?: boolean
+  sanitizePart?: (part: UIMessage['parts'][number]) => UIMessage['parts'][number]
+}
+
+const DEFAULT_DROP_TOOL_STATES = ['input-streaming']
+
+function shouldDropStepPart(type: string) {
+  // `toUIMessageStreamResponse({ sendReasoning: true })` can emit step boundary parts which are
+  // useful for streaming UX, but not useful (and sometimes invalid) when re-hydrated from storage.
+  return type === 'step-start' || type === 'step-finish'
+}
+
+function normalizeParts(
+  parts: Array<unknown>,
+  options: Required<
+    Pick<NormalizeUiMessagesOptions, 'dropReasoningParts' | 'dropStepParts' | 'dropNonTextParts'>
+  > & {
+    dropToolStates: string[]
+    sanitizePart?: NormalizeUiMessagesOptions['sanitizePart']
+  }
+): UIMessage['parts'] {
+  const normalized: UIMessage['parts'] = []
+
+  for (const rawPart of parts) {
+    if (!isRecord(rawPart)) continue
+    const type = typeof rawPart.type === 'string' ? rawPart.type : undefined
+    if (type == null) continue
+
+    if (options.dropStepParts && shouldDropStepPart(type)) continue
+
+    if (type === 'text') {
+      if (typeof rawPart.text !== 'string') continue
+      normalized.push(rawPart as any)
+      continue
+    }
+
+    if (options.dropNonTextParts) continue
+
+    if (type === 'reasoning') {
+      if (options.dropReasoningParts) continue
+      if (typeof rawPart.text !== 'string') continue
+      normalized.push(rawPart as any)
+      continue
+    }
+
+    if (type.startsWith('tool-')) {
+      const state = typeof rawPart.state === 'string' ? rawPart.state : undefined
+      if (state == null) continue
+      if (typeof rawPart.toolCallId !== 'string') continue
+      if (!('input' in rawPart)) continue
+      if (options.dropToolStates.includes(state)) continue
+      if (state === 'output-available' && !('output' in rawPart)) continue
+      if (state === 'output-error' && typeof rawPart.errorText !== 'string') continue
+      if (options.sanitizePart) {
+        normalized.push(options.sanitizePart(rawPart as any))
+      } else {
+        normalized.push(rawPart as any)
+      }
+      continue
+    }
+
+    if (type === 'dynamic-tool') {
+      const state = typeof rawPart.state === 'string' ? rawPart.state : undefined
+      if (state == null) continue
+      if (typeof rawPart.toolName !== 'string') continue
+      if (typeof rawPart.toolCallId !== 'string') continue
+      if (!('input' in rawPart)) continue
+      if (options.dropToolStates.includes(state)) continue
+      if (state === 'output-available' && !('output' in rawPart)) continue
+      if (state === 'output-error' && typeof rawPart.errorText !== 'string') continue
+      if (options.sanitizePart) {
+        normalized.push(options.sanitizePart(rawPart as any))
+      } else {
+        normalized.push(rawPart as any)
+      }
+      continue
+    }
+
+    // For anything else (data parts, sources, etc.), keep it and let `validateUIMessages`
+    // decide whether it matches the current AI SDK schemas.
+    normalized.push(rawPart as any)
+  }
+
+  return normalized
+}
+
+/**
+ * Normalizes message history read from storage into a shape the AI SDK can reliably parse.
+ * This is intentionally strict: unknown/invalid parts are dropped to avoid schema errors.
+ */
+export function normalizeUiMessages(
+  messages: Array<unknown>,
+  {
+    dropReasoningParts = false,
+    dropStepParts = true,
+    dropEmptyMessages = true,
+    dropToolStates = DEFAULT_DROP_TOOL_STATES,
+    dropNonTextParts = false,
+    sanitizePart,
+  }: NormalizeUiMessagesOptions = {}
+): UIMessage[] {
+  const out: UIMessage[] = []
+
+  for (let i = 0; i < messages.length; i++) {
+    const rawMessage = messages[i]
+    if (!isRecord(rawMessage)) continue
+
+    const role = typeof rawMessage.role === 'string' ? rawMessage.role : undefined
+    if (
+      role !== 'user' &&
+      role !== 'assistant' &&
+      role !== 'system' &&
+      role !== 'tool' &&
+      role !== 'data'
+    ) {
+      continue
+    }
+
+    const id = typeof rawMessage.id === 'string' ? rawMessage.id : `msg-${i}`
+
+    const parts =
+      Array.isArray(rawMessage.parts) && rawMessage.parts.length > 0
+        ? rawMessage.parts
+        : typeof rawMessage.content === 'string'
+          ? [{ type: 'text', text: rawMessage.content }]
+          : []
+
+    const normalizedParts = normalizeParts(parts, {
+      dropReasoningParts,
+      dropStepParts,
+      dropToolStates,
+      dropNonTextParts,
+      sanitizePart,
+    })
+
+    if (dropEmptyMessages && normalizedParts.length === 0) continue
+
+    const message: UIMessage = {
+      id,
+      role: role as UIMessage['role'],
+      parts: normalizedParts,
+    }
+
+    if (rawMessage.metadata != null) {
+      message.metadata = rawMessage.metadata as any
+    }
+
+    out.push(message)
+  }
+
+  return out
+}

apps/studio/lib/ai/tools/index.ts

@@ -13,15 +13,17 @@ export const getTools = async ({
   authorization,
   aiOptInLevel,
   accessToken,
+  chatId,
 }: {
   projectRef: string
   connectionString: string
   authorization?: string
   aiOptInLevel: AiOptInLevel
   accessToken?: string
+  chatId?: string
 }) => {
-  // Always include rendering tools
-  let tools: ToolSet = getRenderingTools()
+  // Always include rendering tools with context for server-side execution
+  let tools: ToolSet = getRenderingTools({ projectRef, chatId, authorization })
 
   // If self-hosted, only add fallback tools
   if (!IS_PLATFORM) {

apps/studio/lib/ai/tools/rendering-tools.ts

@@ -1,7 +1,14 @@
 import { tool } from 'ai'
+import { updateChatSession } from 'data/chat-sessions/chat-session-update-mutation'
 import { z } from 'zod'
 
-export const getRenderingTools = () => ({
+type RenderingToolsContext = {
+  projectRef: string
+  chatId?: string
+  authorization?: string
+}
+
+export const getRenderingTools = (context?: RenderingToolsContext) => ({
   execute_sql: tool({
     description: 'Asks the user to execute a SQL statement and return the results',
     inputSchema: z.object({
@@ -34,8 +41,29 @@ export const getRenderingTools = () => ({
     inputSchema: z.object({
       newName: z.string().describe('The new name for the chat session. Five words or less.'),
     }),
-    execute: async () => {
-      return { status: 'Chat request sent to client' }
+    execute: async ({ newName }) => {
+      if (!context?.projectRef || !context?.chatId) {
+        return { status: 'error', message: 'Missing project or chat context' }
+      }
+
+      try {
+        const headers = context.authorization ? { Authorization: context.authorization } : undefined
+        await updateChatSession(
+          {
+            id: context.chatId,
+            projectRef: context.projectRef,
+            name: newName,
+          },
+          headers
+        )
+        return { status: 'success', newName }
+      } catch (error) {
+        console.error('Failed to rename chat:', error)
+        return {
+          status: 'error',
+          message: error instanceof Error ? error.message : 'Failed to rename chat',
+        }
+      }
     },
   }),
 })

apps/studio/lib/api/generate-v4.test.ts

@@ -14,19 +14,14 @@ test('generateV4 calls the tool sanitizer', async () => {
     headers: {
       authorization: 'Bearer test-token',
     },
+    on: vi.fn(),
     body: {
-      messages: [
-        {
-          role: 'assistant',
-          parts: [
-            {
-              type: 'tool-execute_sql',
-              state: 'output-available',
-              output: 'test output',
-            },
-          ],
-        },
-      ],
+      message: {
+        id: 'user-msg-1',
+        role: 'user',
+        parts: [{ type: 'text', text: 'Hello' }],
+      },
+      chatId: '00000000-0000-0000-0000-000000000000',
       projectRef: 'test-project',
       connectionString: 'test-connection',
       orgSlug: 'test-org',
@@ -37,6 +32,9 @@ test('generateV4 calls the tool sanitizer', async () => {
     status: vi.fn(() => mockRes),
     json: vi.fn(() => mockRes),
     setHeader: vi.fn(() => mockRes),
+    writeHead: vi.fn(() => mockRes),
+    write: vi.fn(() => true),
+    end: vi.fn(() => mockRes),
   }
 
   vi.mock('lib/ai/org-ai-details', () => ({
@@ -59,16 +57,68 @@ test('generateV4 calls the tool sanitizer', async () => {
     executeSql: vi.fn().mockResolvedValue({ result: [] }),
   }))
 
+  vi.mock('data/chat-sessions/chat-session-messages-query', () => ({
+    getChatSessionMessages: vi.fn().mockResolvedValue([
+      {
+        id: 'assistant-tool-msg-1',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'tool-execute_sql',
+            toolCallId: 'call-123',
+            state: 'output-available',
+            input: { sql: 'select 1', label: 'Test', chartConfig: { view: 'table' }, isWriteQuery: false },
+            output: [{ ok: true }],
+          },
+        ],
+      },
+    ]),
+  }))
+
+  vi.mock('data/chat-sessions/chat-session-messages-create', () => ({
+    createChatSessionMessages: vi.fn().mockResolvedValue({}),
+  }))
+
   vi.mock('lib/ai/tools', () => ({
     getTools: vi.fn().mockResolvedValue({}),
   }))
 
   vi.mock('ai', () => ({
+    createUIMessageStream: vi.fn().mockImplementation(({ execute, onFinish }) => {
+      const mockWriter = {
+        write: vi.fn(),
+        merge: vi.fn(),
+      }
+      // Execute the stream handler
+      execute({ writer: mockWriter })
+      // Call onFinish with mock response message
+      onFinish?.({ messages: [], responseMessage: null })
+      // Return a mock ReadableStream
+      return new ReadableStream({
+        start(controller) {
+          controller.close()
+        },
+      })
+    }),
+    createUIMessageStreamResponse: vi.fn().mockReturnValue({
+      status: 200,
+      headers: new Headers(),
+      body: null,
+    }),
     streamText: vi.fn().mockReturnValue({
-      pipeUIMessageStreamToResponse: vi.fn(),
+      toUIMessageStream: vi.fn().mockReturnValue(
+        new ReadableStream({
+          start(controller) {
+            controller.close()
+          },
+        })
+      ),
+      consumeStream: vi.fn(),
     }),
     convertToModelMessages: vi.fn((msgs) => msgs),
     stepCountIs: vi.fn(),
+    generateId: vi.fn(() => 'msg-test'),
+    TypeValidationError: class TypeValidationError extends Error {},
   }))
 
   await generateV4(mockReq as any, mockRes as any)

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

@@ -1,7 +1,9 @@
 import pgMeta from '@supabase/pg-meta'
 import {
   convertToModelMessages,
-  createIdGenerator,
+  createUIMessageStream,
+  createUIMessageStreamResponse,
+  generateId,
   type ModelMessage,
   stepCountIs,
   streamText,
@@ -12,23 +14,24 @@ import type { NextApiRequest, NextApiResponse } from 'next'
 import z from 'zod'
 
 import { IS_PLATFORM } from 'common'
-import { get, post } from 'data/fetchers'
+import { createChatSessionMessages } from 'data/chat-sessions/chat-session-messages-create'
+import { getChatSessionMessages } from 'data/chat-sessions/chat-session-messages-query'
 import { executeSql } from 'data/sql/execute-sql-query'
 import type { AiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
+import { cleanMessage } from 'lib/ai/clean-message'
 import { getModel } from 'lib/ai/model'
 import { getOrgAIDetails } from 'lib/ai/org-ai-details'
 import {
   CHAT_PROMPT,
   EDGE_FUNCTION_PROMPT,
   GENERAL_PROMPT,
-  PG_BEST_PRACTICES,
-  RLS_PROMPT,
-  REALTIME_PROMPT,
-  SECURITY_PROMPT,
   LIMITATIONS_PROMPT,
+  PG_BEST_PRACTICES,
+  REALTIME_PROMPT,
+  RLS_PROMPT,
+  SECURITY_PROMPT,
 } from 'lib/ai/prompts'
 import { getTools } from 'lib/ai/tools'
-import { sanitizeMessagePart } from 'lib/ai/tools/tool-sanitizer'
 import apiWrapper from 'lib/api/apiWrapper'
 import { executeQuery } from 'lib/api/self-hosted/query'
 
@@ -103,12 +106,19 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
   let previousMessages: UIMessage[] = []
   if (chatId && projectRef) {
     try {
-      const { data } = await get(`/v1/projects/{ref}/chat-sessions/{id}/messages`, {
-        params: { path: { ref: projectRef, id: chatId } },
+      const data = await getChatSessionMessages({ projectRef, id: chatId }, undefined, {
+        ...(authorization && { Authorization: authorization }),
       })
-      if (data) {
-        previousMessages = data
-      }
+      console.log('Loaded messages from DB:', JSON.stringify(data, null, 2))
+      // Clean loaded messages to match UIMessage format expected by AI SDK
+      previousMessages = (data || []).map((msg: any) => ({
+        id: msg.id,
+        role: msg.role,
+        parts: msg.parts || [],
+        ...(msg.metadata != null && typeof msg.metadata === 'object'
+          ? { metadata: msg.metadata }
+          : {}),
+      })) as UIMessage[]
     } catch (error) {
       console.error('Failed to load previous messages:', error)
       // Continue without previous messages
@@ -142,31 +152,9 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
     }
   }
 
-  // Only returns last 7 messages
-  // Filters out tools with invalid states
-  // Filters out tool outputs based on opt-in level using renderingToolOutputParser
-  const messages = (rawMessages || []).slice(-7).map((msg: any) => {
-    if (msg && msg.role === 'assistant' && 'results' in msg) {
-      const cleanedMsg = { ...msg }
-      delete cleanedMsg.results
-      return cleanedMsg
-    }
-    if (msg && msg.role === 'assistant' && msg.parts) {
-      const cleanedParts = msg.parts
-        .filter((part: any) => {
-          if (part.type.startsWith('tool-')) {
-            const invalidStates = ['input-streaming', 'input-available', 'output-error']
-            return !invalidStates.includes(part.state)
-          }
-          return true
-        })
-        .map((part: any) => {
-          return sanitizeMessagePart(part, aiOptInLevel)
-        })
-      return { ...msg, parts: cleanedParts }
-    }
-    return msg
-  })
+  // Use raw messages directly - take last 7 for context window management
+  const messagesForModel = rawMessages.slice(-7)
+  const originalMessages = rawMessages
 
   const {
     model,
@@ -225,7 +213,7 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
 
     // Note: these must be of type `CoreMessage` to prevent AI SDK from stripping `providerOptions`
     // https://github.com/vercel/ai/blob/81ef2511311e8af34d75e37fc8204a82e775e8c3/packages/ai/core/prompt/standardize-prompt.ts#L83-L88
-    const coreMessages: ModelMessage[] = [
+    const coreMessagesBase: ModelMessage[] = [
       {
         role: 'system',
         content: system,
@@ -238,45 +226,107 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
         // Add any dynamic context here
         content: `The user's current project is ${projectRef}. Their available schemas are: ${schemasString}. The current chat name is: ${chatName}`,
       },
-      ...convertToModelMessages(messages),
     ]
 
     const abortController = new AbortController()
     req.on('close', () => abortController.abort())
     req.on('aborted', () => abortController.abort())
 
-    // Get tools
+    // Get tools with chat context for server-side tool execution
     const tools = await getTools({
       projectRef,
       connectionString,
       authorization,
       aiOptInLevel,
       accessToken,
+      chatId,
     })
 
-    const result = streamText({
-      model,
-      stopWhen: stepCountIs(5),
-      messages: coreMessages,
-      ...(providerOptions && { providerOptions }),
-      tools,
-      abortSignal: abortController.signal,
-    })
+    // Skip validation for now - just use raw messages
+    // The validateUIMessages is too strict about tool part schemas
+    const validatedMessagesForModel = messagesForModel
+    const validatedOriginalMessages = originalMessages
 
-    const uiMessageStreamResponse = result.toUIMessageStreamResponse({
-      originalMessages: rawMessages,
-      generateMessageId: createIdGenerator({ prefix: 'msg', size: 16 }),
-      sendReasoning: true,
-      onFinish: async ({ messages }) => {
-        // Save messages to database if chatId is provided
-        if (chatId && projectRef) {
+    console.log('Messages for model:', JSON.stringify(validatedMessagesForModel, null, 2))
+
+    let coreMessages: ModelMessage[]
+    try {
+      coreMessages = [...coreMessagesBase, ...convertToModelMessages(validatedMessagesForModel)]
+    } catch (error) {
+      console.error('convertToModelMessages failed:', error)
+      // Fall back to text-only messages if conversion fails
+      const textOnlyMessages = validatedMessagesForModel.map((msg) => ({
+        id: msg.id,
+        role: msg.role,
+        parts: msg.parts.filter((part: { type: string }) => part.type === 'text'),
+      }))
+      console.log('Falling back to text-only messages:', JSON.stringify(textOnlyMessages, null, 2))
+      coreMessages = [
+        ...coreMessagesBase,
+        ...convertToModelMessages(textOnlyMessages as UIMessage[]),
+      ]
+    }
+
+    // Headers for API calls
+    const apiHeaders = {
+      'Content-Type': 'application/json',
+      ...(authorization && { Authorization: authorization }),
+    }
+
+    // Create UI message stream with immediate persistence
+    const stream = createUIMessageStream({
+      execute: async ({ writer }) => {
+        // 1. Save user message BEFORE streaming starts
+        const userMessage = message as UIMessage
+        if (chatId && projectRef && userMessage.role === 'user') {
           try {
-            await post(`/v1/projects/{ref}/chat-sessions/{id}/messages`, {
-              params: { path: { ref: projectRef, id: chatId } },
-              body: { messages },
-            })
+            console.log('Saving user message:', JSON.stringify(userMessage, null, 2))
+            await createChatSessionMessages(
+              { projectRef, id: chatId, messages: [userMessage] },
+              apiHeaders
+            )
           } catch (error) {
-            console.error('Failed to save messages:', error)
+            console.error('Failed to save user message:', error)
+            // Continue even if save fails - don't block the streaming
+          }
+        }
+
+        // 2. Write start events for assistant response (only for user messages)
+        if (userMessage.role === 'user') {
+          writer.write({ type: 'start', messageId: generateId() })
+          writer.write({ type: 'start-step' })
+        }
+
+        // 3. Execute streamText and merge output
+        const result = streamText({
+          model,
+          stopWhen: stepCountIs(5),
+          messages: coreMessages,
+          ...(providerOptions && { providerOptions }),
+          tools,
+          abortSignal: abortController.signal,
+        })
+
+        // Consume the stream to ensure it runs
+        result.consumeStream()
+
+        // Merge the stream output (sendStart: false since we already wrote start events)
+        writer.merge(result.toUIMessageStream({ sendStart: false, sendReasoning: true }))
+      },
+      originalMessages: validatedOriginalMessages,
+      onFinish: async ({ responseMessage }) => {
+        // 4. Save assistant message after streaming completes
+        console.log('responseMessage:', JSON.stringify(responseMessage))
+        if (chatId && projectRef && responseMessage) {
+          try {
+            const cleanedMessage = cleanMessage(responseMessage)
+            console.log('Saving assistant message:', JSON.stringify(cleanedMessage, null, 2))
+            await createChatSessionMessages(
+              { projectRef, id: chatId, messages: [cleanedMessage] },
+              apiHeaders
+            )
+          } catch (error) {
+            console.error('Failed to save assistant message:', error)
           }
         }
       },
@@ -297,6 +347,9 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
       },
     })
 
+    // Create response from stream
+    const uiMessageStreamResponse = createUIMessageStreamResponse({ stream })
+
     // Stream the response
     res.writeHead(uiMessageStreamResponse.status, {
       'Content-Type': 'text/plain; charset=utf-8',