From 0f19e48a0aee0b7cbafc8cc0cdcd3ad23b48998a Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Wed, 17 Dec 2025 12:34:37 +1000 Subject: [PATCH] fixes --- .../ui/AIAssistantPanel/AIAssistant.tsx | 158 +++++++++++----- .../chat-session-message-delete-mutation.ts | 47 +++++ .../chat-session-message-delete.ts | 34 ++++ .../chat-session-messages-create-mutation.ts | 56 ++++++ .../chat-session-messages-create.ts | 31 ++++ .../chat-session-messages-query.ts | 20 +- .../chat-session-update-mutation.ts | 8 +- .../data/chat-sessions/chat-sessions-query.ts | 10 +- apps/studio/lib/ai/clean-message.ts | 39 ++++ .../lib/ai/normalize-ui-messages.test.ts | 77 ++++++++ apps/studio/lib/ai/normalize-ui-messages.ts | 165 +++++++++++++++++ apps/studio/lib/ai/tools/index.ts | 6 +- apps/studio/lib/ai/tools/rendering-tools.ts | 34 +++- apps/studio/lib/api/generate-v4.test.ts | 76 ++++++-- apps/studio/pages/api/ai/sql/generate-v4.ts | 173 ++++++++++++------ 15 files changed, 803 insertions(+), 131 deletions(-) create mode 100644 apps/studio/data/chat-sessions/chat-session-message-delete-mutation.ts create mode 100644 apps/studio/data/chat-sessions/chat-session-message-delete.ts create mode 100644 apps/studio/data/chat-sessions/chat-session-messages-create-mutation.ts create mode 100644 apps/studio/data/chat-sessions/chat-session-messages-create.ts create mode 100644 apps/studio/lib/ai/clean-message.ts create mode 100644 apps/studio/lib/ai/normalize-ui-messages.test.ts create mode 100644 apps/studio/lib/ai/normalize-ui-messages.ts diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index 7e967278858..e971dbcfb4d 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/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(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( { 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(null) + const hasHydratedChatMessagesRef = useRef(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} diff --git a/apps/studio/data/chat-sessions/chat-session-message-delete-mutation.ts b/apps/studio/data/chat-sessions/chat-session-message-delete-mutation.ts new file mode 100644 index 00000000000..f0e5cc8bd6b --- /dev/null +++ b/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, + }) +} diff --git a/apps/studio/data/chat-sessions/chat-session-message-delete.ts b/apps/studio/data/chat-sessions/chat-session-message-delete.ts new file mode 100644 index 00000000000..bed3fad283c --- /dev/null +++ b/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> +export type ChatSessionMessageDeleteError = ResponseError diff --git a/apps/studio/data/chat-sessions/chat-session-messages-create-mutation.ts b/apps/studio/data/chat-sessions/chat-session-messages-create-mutation.ts new file mode 100644 index 00000000000..50c8604bf71 --- /dev/null +++ b/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, + }) +} + diff --git a/apps/studio/data/chat-sessions/chat-session-messages-create.ts b/apps/studio/data/chat-sessions/chat-session-messages-create.ts new file mode 100644 index 00000000000..3556e03ca18 --- /dev/null +++ b/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> +export type ChatSessionMessagesCreateError = ResponseError diff --git a/apps/studio/data/chat-sessions/chat-session-messages-query.ts b/apps/studio/data/chat-sessions/chat-session-messages-query.ts index 83537932575..da2dab81db5 100644 --- a/apps/studio/data/chat-sessions/chat-session-messages-query.ts +++ b/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> diff --git a/apps/studio/data/chat-sessions/chat-session-update-mutation.ts b/apps/studio/data/chat-sessions/chat-session-update-mutation.ts index cad560d5b2c..3bf748e245b 100644 --- a/apps/studio/data/chat-sessions/chat-session-update-mutation.ts +++ b/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) diff --git a/apps/studio/data/chat-sessions/chat-sessions-query.ts b/apps/studio/data/chat-sessions/chat-sessions-query.ts index 00baa1bd904..c26e717f8ea 100644 --- a/apps/studio/data/chat-sessions/chat-sessions-query.ts +++ b/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> diff --git a/apps/studio/lib/ai/clean-message.ts b/apps/studio/lib/ai/clean-message.ts new file mode 100644 index 00000000000..16b2cd62903 --- /dev/null +++ b/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) +} diff --git a/apps/studio/lib/ai/normalize-ui-messages.test.ts b/apps/studio/lib/ai/normalize-ui-messages.test.ts new file mode 100644 index 00000000000..113bb5cc790 --- /dev/null +++ b/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...' }) + }) +}) diff --git a/apps/studio/lib/ai/normalize-ui-messages.ts b/apps/studio/lib/ai/normalize-ui-messages.ts new file mode 100644 index 00000000000..5427048aca2 --- /dev/null +++ b/apps/studio/lib/ai/normalize-ui-messages.ts @@ -0,0 +1,165 @@ +import type { UIMessage } from 'ai' + +type UnknownRecord = Record + +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, + options: Required< + Pick + > & { + 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, + { + 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 +} diff --git a/apps/studio/lib/ai/tools/index.ts b/apps/studio/lib/ai/tools/index.ts index 9082cfdb81f..a1db6bed27d 100644 --- a/apps/studio/lib/ai/tools/index.ts +++ b/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) { diff --git a/apps/studio/lib/ai/tools/rendering-tools.ts b/apps/studio/lib/ai/tools/rendering-tools.ts index 4f5e9704ff4..aca63671162 100644 --- a/apps/studio/lib/ai/tools/rendering-tools.ts +++ b/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', + } + } }, }), }) diff --git a/apps/studio/lib/api/generate-v4.test.ts b/apps/studio/lib/api/generate-v4.test.ts index a04be6ef01a..18ebc0d4d47 100644 --- a/apps/studio/lib/api/generate-v4.test.ts +++ b/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) diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index f1b3187ca9c..d0fd21b960d 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/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',