diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index 02c2844a139..8b2b0b1776f 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx @@ -1,59 +1,14 @@ import type { UIMessage as MessageType } from '@ai-sdk/react' -import { useChat } from '@ai-sdk/react' -import { lastAssistantMessageIsCompleteWithApprovalResponses } from 'ai' -import { LOCAL_STORAGE_KEYS, useFlag } from 'common' -import { useParams, useSearchParamsShallow } from 'common/hooks' -import { AnimatePresence, motion } from 'framer-motion' -import { Eraser, Pencil, X } from 'lucide-react' +import { useParams } from 'common/hooks' import { useRouter } from 'next/router' -import { Fragment, useCallback, useEffect, useMemo, useRef, useState } from 'react' -import { Button, cn, KeyboardShortcut } from 'ui' -import { Admonition } from 'ui-patterns/Admonition' +import { useEffect } from 'react' -import { AlertError } from '../AlertError' -import { ButtonTooltip } from '../ButtonTooltip' -import { ErrorBoundary } from '../ErrorBoundary/ErrorBoundary' -import { InlineLinkClassName } from '../InlineLink' -import { ASSISTANT_ERRORS } from './AiAssistant.constants' -import { - containsLogsSnippets, - hasPendingToolApproval, - onErrorChat, - resolvePendingToolApprovalsAsDenied, -} from './AIAssistant.utils' import { AIAssistantHeader } from './AIAssistantHeader' -import { AIOnboarding } from './AIOnboarding' -import { AssistantChatForm } from './AssistantChatForm' -import { - Conversation, - ConversationContent, - ConversationScrollButton, -} from './elements/Conversation' -import { Message } from './Message' -import { Markdown } from '@/components/interfaces/Markdown' +import { AssistantChat } from './AssistantChat' import { resolveSnippetSource } from '@/components/interfaces/SQLEditor/querySource' import { SIDEBAR_KEYS } from '@/components/layouts/ProjectLayout/LayoutSidebar/LayoutSidebarProvider' -import { useCheckOpenAIKeyQuery } from '@/data/ai/check-api-key-query' -import { useRateMessageMutation } from '@/data/ai/rate-message-mutation' -import { useTablesQuery } from '@/data/tables/tables-query' -import { useCheckEntitlements } from '@/hooks/misc/useCheckEntitlements' -import { useLocalStorageQuery } from '@/hooks/misc/useLocalStorage' -import { useOrgAiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' -import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization' -import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject' -import type { AssistantMessageMetadata } from '@/lib/ai/assistant-message-metadata' -import { getParallelApprovalIdsToReject } from '@/lib/ai/message-utils' -import { - DEFAULT_ASSISTANT_BASE_MODEL_ID, - defaultAssistantModelId, - isAssistantBaseModelId, - isKnownAssistantModelId, -} from '@/lib/ai/model.utils' -import { IS_PLATFORM } from '@/lib/constants' -import { uuidv4 } from '@/lib/helpers' -import { useTrack } from '@/lib/telemetry/track' -import type { AssistantModel, SqlSnippet } from '@/state/ai-assistant-state' import { useAiAssistantState, useAiAssistantStateSnapshot } from '@/state/ai-assistant-state' +import type { SqlSnippet } from '@/state/ai-assistant-state' import { SHORTCUT_IDS } from '@/state/shortcuts/registry' import { useShortcut } from '@/state/shortcuts/useShortcut' import { useSidebarManagerSnapshot } from '@/state/sidebar-manager-state' @@ -64,646 +19,89 @@ interface AIAssistantProps { className?: string } +type CurrentQuerySnippet = Exclude + +const isSameSnippet = (snippet: SqlSnippet, currentQuery: CurrentQuerySnippet) => + typeof snippet !== 'string' && + snippet.label === currentQuery.label && + snippet.content === currentQuery.content && + snippet.source === currentQuery.source + export const AIAssistant = ({ className }: AIAssistantProps) => { const router = useRouter() const { id: entityId, source: sourceParam } = useParams() - const { data: project } = useSelectedProjectQuery() - const searchParams = useSearchParamsShallow() - - const { data: selectedOrganization, isPending: isLoadingOrganization } = - useSelectedOrganizationQuery() - - useShortcut(SHORTCUT_IDS.AI_ASSISTANT_CANCEL_EDIT, () => cancelEdit()) - useShortcut(SHORTCUT_IDS.AI_ASSISTANT_NEW_CHAT, () => snap.newChat()) - - const disablePrompts = useFlag('disableAssistantPrompts') - const { snippets } = useSqlEditorV2StateSnapshot() const snap = useAiAssistantStateSnapshot() const state = useAiAssistantState() + const { snippets } = useSqlEditorV2StateSnapshot() const { activeSidebar, closeSidebar } = useSidebarManagerSnapshot() + const shortcutsEnabled = activeSidebar?.id === SIDEBAR_KEYS.AI_ASSISTANT - const { hasAccess: hasAccessToAdvanceModel, isLoading: isLoadingEntitlements } = - useCheckEntitlements('assistant.advance_model') + const handleNewChat = () => state.newChat() - const selectedModel = useMemo(() => { - // While entitlements are loading, use the stored model without enforcing access - if (isLoadingEntitlements) { - return snap.model ?? DEFAULT_ASSISTANT_BASE_MODEL_ID - } - - const defaultModel = defaultAssistantModelId(hasAccessToAdvanceModel) - const model = snap.model ?? defaultModel - - if (!isKnownAssistantModelId(model)) return defaultModel - if (!hasAccessToAdvanceModel && !isAssistantBaseModelId(model)) { - return DEFAULT_ASSISTANT_BASE_MODEL_ID - } - - return model - }, [isLoadingEntitlements, hasAccessToAdvanceModel, snap.model]) - - const [updatedOptInSinceMCP] = useLocalStorageQuery( - LOCAL_STORAGE_KEYS.AI_ASSISTANT_MCP_OPT_IN, - false - ) - - const inputRef = useRef(null) - - const { aiOptInLevel, isHipaaProjectDisallowed } = useOrgAiOptInLevel() - // Whether attached queries are sent at all. One definition, shared by the chat form - // (which folds them into the message text) and the message metadata (which states - // whether any of them was a logs query), so the two can't disagree. - const includeSnippetsInMessage = aiOptInLevel !== 'disabled' - const showMetadataWarning = - IS_PLATFORM && - !!selectedOrganization && - (aiOptInLevel === 'disabled' || aiOptInLevel === 'schema') - - // Add a ref to store the last user message - const lastUserMessageRef = useRef(null) - - // Keep latest selected organization to avoid stale values in useChat transport - const selectedOrganizationRef = useRef(selectedOrganization) - useEffect(() => { - selectedOrganizationRef.current = selectedOrganization - }, [selectedOrganization]) - - const [value, setValue] = useState(snap.initialInput || '') - const [editingMessageId, setEditingMessageId] = useState(null) - const [isResubmitting, setIsResubmitting] = useState(false) - const [messageRatings, setMessageRatings] = useState>({}) - - const { data: check, isSuccess } = useCheckOpenAIKeyQuery() - const isApiKeySet = !!check?.hasKey - - const { mutateAsync: rateMessage } = useRateMessageMutation() + useShortcut(SHORTCUT_IDS.AI_ASSISTANT_NEW_CHAT, handleNewChat, { + enabled: shortcutsEnabled, + }) const isInSQLEditor = router.pathname.includes('/sql/[id]') const snippet = snippets[entityId ?? ''] const snippetContent = snippet?.snippet?.content?.unchecked_sql - const openSnippetSource = isInSQLEditor ? resolveSnippetSource(snippet?.snippet, sourceParam) : undefined - const { data: tables } = useTablesQuery( - { - projectRef: project?.ref, - connectionString: project?.connectionString, - schema: 'public', - }, - { enabled: isApiKeySet } - ) - - const currentTable = tables?.find((t) => t.id.toString() === entityId) - const currentSchema = searchParams?.get('schema') ?? 'public' - - // Update context in state useEffect(() => { - state.setContext({ - projectRef: project?.ref, - orgSlug: selectedOrganizationRef.current?.slug, - connectionString: project?.connectionString ?? '', - }) - }, [project?.ref, project?.connectionString, selectedOrganizationRef.current?.slug, state]) + if (!shortcutsEnabled || !isInSQLEditor || !snippetContent) return - const track = useTrack() + const currentQuery = { + label: 'Current Query', + content: snippetContent, + source: openSnippetSource, + } + state.setSqlSnippets([currentQuery]) - const { - messages: chatMessages, - status: chatStatus, - error, - sendMessage, - setMessages, - addToolApprovalResponse, - stop, - regenerate, - } = useChat({ - id: snap.activeChatId, - ...(snap.activeChatId && snap.chatInstances[snap.activeChatId] - ? { chat: snap.chatInstances[snap.activeChatId] } - : {}), - sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses, - onError: onErrorChat, - }) + return () => { + const currentSnippets = state.sqlSnippets + const remainingSnippets = currentSnippets?.filter( + (snippet) => !isSameSnippet(snippet, currentQuery) + ) - const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming' - const hasPendingApproval = hasPendingToolApproval(chatMessages) - const supportMetadata = snap.activeChat?.supportMetadata - const isSupportChat = !!supportMetadata?.isSupportChat - const isSupportChatClosed = isSupportChat && supportMetadata.lifecycleStatus !== 'bot_active' - const activeChatId = snap.activeChatId - const supportConversationId = supportMetadata?.frontConversationId - const isChatInputDisabled = - !isApiKeySet || disablePrompts || isLoadingOrganization || isSupportChatClosed - - const branchedFrom = snap.activeChat?.branchedFrom - const branchedConversation = branchedFrom ? snap.chats[branchedFrom.chatId] : undefined - - const deleteMessageFromHere = useCallback( - (messageId: string) => { - // Find the message index in current chatMessages - const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId) - if (messageIndex === -1) return - - if (isChatLoading) stop() - - snap.deleteMessagesAfter(messageId, { includeSelf: true }) - - const updatedMessages = chatMessages.slice(0, messageIndex) - setMessages(updatedMessages) - }, - [snap, setMessages, chatMessages, isChatLoading, stop] - ) - - const editMessage = useCallback( - (messageId: string) => { - const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId) - if (messageIndex === -1) return - - // Target message - const messageToEdit = chatMessages[messageIndex] - - // Activate editing mode - setEditingMessageId(messageId) - const textContent = - messageToEdit.parts - ?.filter((part) => part.type === 'text') - .map((part) => part.text) - .join('') ?? '' - setValue(textContent) - - setTimeout(() => { - if (inputRef.current) { - inputRef?.current?.focus() - - // [Joshen] This is just to make the cursor go to the end of the text when focusing - const val = inputRef.current.value - inputRef.current.value = '' - inputRef.current.value = val - } - }, 100) - }, - [chatMessages, setValue] - ) - - const cancelEdit = useCallback(() => { - setEditingMessageId(null) - setValue('') - }, [setValue]) - - const handleRateMessage = useCallback( - async (messageId: string, rating: 'positive' | 'negative', reason?: string) => { - if (!project?.ref || !selectedOrganization?.slug) return - - // Optimistically update UI - setMessageRatings((prev) => ({ ...prev, [messageId]: rating })) - - try { - const result = await rateMessage({ - rating, - messages: chatMessages, - messageId, - projectRef: project.ref, - orgSlug: selectedOrganization.slug, - reason, - spanId: state.messageSpanIds[messageId], - }) - - track('assistant_message_rating_submitted', { - rating, - category: result.category, - ...(reason && { reason }), - chatId: state.activeChatId, - }) - } catch (error) { - console.error('Failed to rate message:', error) - // Rollback on error - setMessageRatings((prev) => { - const { [messageId]: _, ...rest } = prev - return rest - }) + if (currentSnippets && remainingSnippets?.length !== currentSnippets.length) { + state.setSqlSnippets(remainingSnippets ?? []) } - }, - [chatMessages, project?.ref, selectedOrganization?.slug, rateMessage, track, state] - ) - - const isContextExceededError = - error && - (error.message?.includes('context_length_exceeded') || - error.message?.includes('exceeds the context window')) - - const renderedMessages = useMemo( - () => - chatMessages.map((message, index) => { - const isBeingEdited = editingMessageId === message.id - const isAfterEditedMessage = editingMessageId - ? chatMessages.findIndex((m) => m.id === editingMessageId) < index - : false - const isLastMessage = index === chatMessages.length - 1 - - return ( - - - {branchedConversation && branchedFrom?.messageId === message.id && ( -
-
-
- Branched from - -
-
-
- )} - - ) - }), - [ - chatMessages, - deleteMessageFromHere, - editMessage, - cancelEdit, - editingMessageId, - chatStatus, - addToolApprovalResponse, - handleRateMessage, - messageRatings, - branchedConversation, - branchedFrom, - snap, - ] - ) - - const hasMessages = chatMessages.length > 0 - - const sendMessageToAssistant = (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) } + }, [shortcutsEnabled, isInSQLEditor, snippetContent, openSnippetSource, state]) - // Read off the attachments this message actually carries, so detaching the - // "Current Query" chip also drops the claim. Gated on the same condition that - // decides whether attachments make it into the text at all: with AI opt-in - // disabled the chip is shown but no query is sent, and claiming otherwise would - // have the server prepend ClickHouse context for a message holding no query. - // Rides on the message rather than the request, so a Retry reproduces the context - // the message was asked in. - const metadata: AssistantMessageMetadata = { - containsLogsSnippets: includeSnippetsInMessage && containsLogsSnippets(snap.sqlSnippets), - } - - const payload = { - role: 'user', - createdAt: new Date(), - parts: [{ type: 'text', text: finalContent }], - id: uuidv4(), - metadata, - } as MessageType - - snap.clearSqlSnippets() - lastUserMessageRef.current = payload - if (hasPendingApproval && !editingMessageId) { - setMessages(resolvePendingToolApprovalsAsDenied(chatMessages)) - } - sendMessage(payload, { - body: { - schema: currentSchema, - table: currentTable?.name, - }, - }) - setValue('') - - if (finalContent.includes('Help me to debug')) { - track('assistant_debug_submitted', { chatId: snap.activeChatId }) - } else { - track('assistant_prompt_submitted', { chatId: snap.activeChatId }) - } - } - - const handleClearMessages = () => { - if (isChatLoading) stop() - snap.clearMessages() - setMessages([]) - lastUserMessageRef.current = null - setEditingMessageId(null) - } - - useEffect(() => { - // Keep "Thinking" visible while stopping and resubmitting during edit - // Only clear once the new response actually starts streaming (or errors) - if (isResubmitting && (chatStatus === 'streaming' || !!error)) { - setIsResubmitting(false) - } - }, [isResubmitting, chatStatus, error]) - - useEffect(() => { - // Approval-required tools can't run in parallel. Auto-deny extras so the model reissues them sequentially. - for (const id of getParallelApprovalIdsToReject(chatMessages)) { - addToolApprovalResponse?.({ - id, - approved: false, - reason: - 'Only one approval-required tool call is allowed per turn. Please reissue this tool call after the current one completes.', - }) - } - }, [chatMessages, addToolApprovalResponse]) - - useEffect(() => { - setValue(snap.initialInput || '') - if (inputRef.current && snap.initialInput) { - inputRef.current.focus() - inputRef.current.setSelectionRange(snap.initialInput.length, snap.initialInput.length) - } - }, [snap.initialInput]) - - useEffect(() => { - const isOpen = activeSidebar?.id === SIDEBAR_KEYS.AI_ASSISTANT - if (isOpen && isInSQLEditor && !!snippetContent) { - snap.setSqlSnippets([ - { label: 'Current Query', content: snippetContent, source: openSnippetSource }, - ]) - } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [activeSidebar?.id, isInSQLEditor, snippetContent, openSnippetSource]) + if (!snap.activeChatId) return null return ( - state.selectChat(chatId)} + onBranchChat={(messageId) => state.branchChat(messageId)} + composerContext={{ + initialInput: snap.initialInput, + sqlSnippets: snap.sqlSnippets as SqlSnippet[] | undefined, + suggestions: snap.suggestions + ? { + title: snap.suggestions.title, + prompts: snap.suggestions.prompts?.map((prompt) => ({ ...prompt })), + } + : undefined, + onSetSqlSnippets: state.setSqlSnippets, + onClearSqlSnippets: state.clearSqlSnippets, }} - actions={[ - { - label: 'Clear messages and refresh', - onClick: () => { - handleClearMessages() - window.location.reload() - }, - }, - ]} - > -
+ renderHeader={(props) => ( closeSidebar(SIDEBAR_KEYS.AI_ASSISTANT)} - showMetadataWarning={showMetadataWarning} - updatedOptInSinceMCP={updatedOptInSinceMCP} - isHipaaProjectDisallowed={isHipaaProjectDisallowed} - aiOptInLevel={aiOptInLevel} /> - {hasMessages ? ( - - - {renderedMessages} - {error && ( - <> - - {isContextExceededError ? ( - - ) : ( - <> - - } - tooltip={{ content: { side: 'bottom', text: 'Clear messages' } }} - /> - - )} -
- } - /> - - )} - {isChatLoading && ( - - )} - -

- The Assistant can make mistakes. Double check responses. -

- - - - ) : ( - setValue(val)} - onFocusInput={() => inputRef.current?.focus()} - /> - )} - - - {editingMessageId && ( - -
- -
-
- - Editing message -
- } - onClick={cancelEdit} - className="w-6 h-6 p-0" - title="Cancel editing" - aria-label="Cancel editing" - tooltip={{ - content: { side: 'top', text: }, - }} - /> -
- -
- - )} - - -
- {isSupportChat && !isSupportChatClosed && ( -
-
-
- - -
-
- )} - - {disablePrompts && ( - - )} - - {isSuccess && !isApiKeySet && ( - - } - /> - )} - - form>textarea]:text-base [&>form>textarea]:md:text-sm [&>form>textarea]:border', - '[&>form>textarea]:rounded-md [&>form>textarea]:outline-hidden!', - '[&>form>textarea]:ring-offset-0! [&>form>textarea]:ring-0!' - )} - loading={isChatLoading} - isEditing={!!editingMessageId} - disabled={isChatInputDisabled} - placeholder={ - hasMessages - ? isSupportChat - ? 'Share details so the assistant can help with your support request...' - : 'Ask a follow up question...' - : (snap.sqlSnippets ?? [])?.length > 0 - ? 'Ask a question or make a change...' - : isSupportChat - ? 'Describe your support issue...' - : 'Chat to Postgres...' - } - value={value} - onValueChange={(e) => setValue(e.target.value)} - onSubmit={(finalMessage) => { - sendMessageToAssistant(finalMessage) - }} - onStop={() => { - stop() - // to save partial responses from the AI - const lastMessage = chatMessages[chatMessages.length - 1] - if (lastMessage && lastMessage.role === 'assistant') { - state.updateMessage(lastMessage) - } - }} - sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined} - onRemoveSnippet={(index) => { - const newSnippets = [...(snap.sqlSnippets ?? [])] - newSnippets.splice(index, 1) - snap.setSqlSnippets(newSnippets) - }} - includeSnippetsInMessage={includeSnippetsInMessage} - selectedModel={selectedModel} - onSelectModel={(model) => snap.setModel(model)} - /> -
-
- + )} + /> ) } diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx index e5631c81b38..11b87ad7653 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistantChatSelector.tsx @@ -24,9 +24,13 @@ import { useShortcut } from '@/state/shortcuts/useShortcut' interface AIAssistantChatSelectorProps { disabled?: boolean + shortcutsEnabled?: boolean } -export const AIAssistantChatSelector = ({ disabled = false }: AIAssistantChatSelectorProps) => { +export const AIAssistantChatSelector = ({ + disabled = false, + shortcutsEnabled = true, +}: AIAssistantChatSelectorProps) => { const snap = useAiAssistantStateSnapshot() const [chatSelectorOpen, setChatSelectorOpen] = useState(false) @@ -35,7 +39,13 @@ export const AIAssistantChatSelector = ({ disabled = false }: AIAssistantChatSel const chats = Object.entries(snap.chats) - useShortcut(SHORTCUT_IDS.AI_ASSISTANT_TOGGLE_HISTORY, () => setChatSelectorOpen((prev) => !prev)) + useShortcut( + SHORTCUT_IDS.AI_ASSISTANT_TOGGLE_HISTORY, + () => setChatSelectorOpen((prev) => !prev), + { + enabled: shortcutsEnabled, + } + ) const handleSelectChat = (id: string) => { snap.selectChat(id) diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistantHeader.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistantHeader.tsx index 49351c8a205..c2f7294590a 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistantHeader.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistantHeader.tsx @@ -33,6 +33,7 @@ import { useSidebarManagerSnapshot } from '@/state/sidebar-manager-state' interface AIAssistantHeaderProps { isChatLoading: boolean + shortcutsEnabled?: boolean onNewChat: () => void onCloseAssistant: () => void showMetadataWarning: boolean @@ -43,6 +44,7 @@ interface AIAssistantHeaderProps { export const AIAssistantHeader = ({ isChatLoading, + shortcutsEnabled = true, onNewChat, onCloseAssistant, showMetadataWarning, @@ -87,15 +89,15 @@ export const AIAssistantHeader = ({ } useShortcut(SHORTCUT_IDS.AI_ASSISTANT_COPY_CHAT_ID, handleCopyChatId, { - enabled: !isChatLoading, + enabled: shortcutsEnabled && !isChatLoading, }) useShortcut(SHORTCUT_IDS.AI_ASSISTANT_OPEN_PERMISSIONS, () => setIsOptInModalOpen(true), { - enabled: !isChatLoading, + enabled: shortcutsEnabled && !isChatLoading, }) useShortcut(SHORTCUT_IDS.AI_ASSISTANT_MAXIMIZE, toggleMaximise, { - enabled: !isChatLoading, + enabled: shortcutsEnabled && !isChatLoading, }) return ( @@ -128,7 +130,7 @@ export const AIAssistantHeader = ({
- + void + onClearSqlSnippets?: () => void +} + +interface AssistantChatProps { + className?: string + chatId: string + shortcutsEnabled?: boolean + onNewChat: () => void + onSelectChat: (chatId: string) => void + onBranchChat: (messageId: string) => void + composerContext?: AssistantChatComposerContext + renderHeader?: (props: AssistantChatHeaderProps) => ReactNode +} + +export const AssistantChat = ({ + className, + chatId, + shortcutsEnabled = true, + onNewChat, + onSelectChat, + onBranchChat, + composerContext, + renderHeader, +}: AssistantChatProps) => { + const { id: entityId } = useParams() + const { data: project } = useSelectedProjectQuery() + const searchParams = useSearchParamsShallow() + + const { data: selectedOrganization, isPending: isLoadingOrganization } = + useSelectedOrganizationQuery() + + const disablePrompts = useFlag('disableAssistantPrompts') + const snap = useAiAssistantStateSnapshot() + const state = useAiAssistantState() + const currentChat = snap.chats[chatId] + + useShortcut(SHORTCUT_IDS.AI_ASSISTANT_CANCEL_EDIT, () => cancelEdit(), { + enabled: shortcutsEnabled, + }) + + const { hasAccess: hasAccessToAdvanceModel, isLoading: isLoadingEntitlements } = + useCheckEntitlements('assistant.advance_model') + + const selectedModel = useMemo(() => { + // While entitlements are loading, use the stored model without enforcing access + if (isLoadingEntitlements) { + return snap.model ?? DEFAULT_ASSISTANT_BASE_MODEL_ID + } + + const defaultModel = defaultAssistantModelId(hasAccessToAdvanceModel) + const model = snap.model ?? defaultModel + + if (!isKnownAssistantModelId(model)) return defaultModel + if (!hasAccessToAdvanceModel && !isAssistantBaseModelId(model)) { + return DEFAULT_ASSISTANT_BASE_MODEL_ID + } + + return model + }, [isLoadingEntitlements, hasAccessToAdvanceModel, snap.model]) + + const [updatedOptInSinceMCP] = useLocalStorageQuery( + LOCAL_STORAGE_KEYS.AI_ASSISTANT_MCP_OPT_IN, + false + ) + + const inputRef = useRef(null) + + const { aiOptInLevel, isHipaaProjectDisallowed } = useOrgAiOptInLevel() + // Whether attached queries are sent at all. One definition, shared by the chat form + // (which folds them into the message text) and the message metadata (which states + // whether any of them was a logs query), so the two can't disagree. + const includeSnippetsInMessage = aiOptInLevel !== 'disabled' + const showMetadataWarning = + IS_PLATFORM && + !!selectedOrganization && + (aiOptInLevel === 'disabled' || aiOptInLevel === 'schema') + + // Add a ref to store the last user message + const lastUserMessageRef = useRef(null) + + const [value, setValue] = useState(composerContext?.initialInput || '') + const [editingMessageId, setEditingMessageId] = useState(null) + const [isResubmitting, setIsResubmitting] = useState(false) + const [messageRatings, setMessageRatings] = useState>({}) + + const { data: check, isSuccess } = useCheckOpenAIKeyQuery() + const isApiKeySet = !!check?.hasKey + + const { mutateAsync: rateMessage } = useRateMessageMutation() + + const { data: tables } = useTablesQuery( + { + projectRef: project?.ref, + connectionString: project?.connectionString, + schema: 'public', + }, + { enabled: isApiKeySet } + ) + + const currentTable = tables?.find((t) => t.id.toString() === entityId) + const currentSchema = searchParams?.get('schema') ?? 'public' + + // Update context in state + useEffect(() => { + state.setContext({ + projectRef: project?.ref, + orgSlug: selectedOrganization?.slug, + connectionString: project?.connectionString ?? '', + }) + }, [project?.ref, project?.connectionString, selectedOrganization?.slug, state]) + + const track = useTrack() + + useEffect(() => { + state.ensureChatInstance(chatId) + }, [chatId, state]) + + const chatInstance = snap.chatInstances[chatId] + + const { + messages: chatMessages, + status: chatStatus, + error, + sendMessage, + setMessages, + addToolApprovalResponse, + stop, + regenerate, + } = useChat({ + id: chatId, + ...(chatInstance ? { chat: chatInstance } : {}), + sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses, + onError: onErrorChat, + }) + + const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming' + const hasPendingApproval = hasPendingToolApproval(chatMessages) + const supportMetadata = currentChat?.supportMetadata + const isSupportChat = !!supportMetadata?.isSupportChat + const isSupportChatClosed = isSupportChat && supportMetadata.lifecycleStatus !== 'bot_active' + const supportConversationId = supportMetadata?.frontConversationId + const isChatInputDisabled = + !isApiKeySet || disablePrompts || isLoadingOrganization || isSupportChatClosed + + const branchedFrom = currentChat?.branchedFrom + const branchedConversation = branchedFrom ? snap.chats[branchedFrom.chatId] : undefined + + const deleteMessageFromHere = useCallback( + (messageId: string) => { + // Find the message index in current chatMessages + const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId) + if (messageIndex === -1) return + + if (isChatLoading) stop() + + snap.deleteMessagesAfter(messageId, { includeSelf: true, chatId }) + + const updatedMessages = chatMessages.slice(0, messageIndex) + setMessages(updatedMessages) + }, + [snap, setMessages, chatMessages, isChatLoading, stop, chatId] + ) + + const editMessage = useCallback( + (messageId: string) => { + const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId) + if (messageIndex === -1) return + + // Target message + const messageToEdit = chatMessages[messageIndex] + + // Activate editing mode + setEditingMessageId(messageId) + const textContent = + messageToEdit.parts + ?.filter((part) => part.type === 'text') + .map((part) => part.text) + .join('') ?? '' + setValue(textContent) + + setTimeout(() => { + if (inputRef.current) { + inputRef?.current?.focus() + + // [Joshen] This is just to make the cursor go to the end of the text when focusing + const val = inputRef.current.value + inputRef.current.value = '' + inputRef.current.value = val + } + }, 100) + }, + [chatMessages, setValue] + ) + + const cancelEdit = useCallback(() => { + setEditingMessageId(null) + setValue('') + }, [setValue]) + + const handleRateMessage = useCallback( + async (messageId: string, rating: 'positive' | 'negative', reason?: string) => { + if (!project?.ref || !selectedOrganization?.slug) return + + // Optimistically update UI + setMessageRatings((prev) => ({ ...prev, [messageId]: rating })) + + try { + const result = await rateMessage({ + rating, + messages: chatMessages, + messageId, + projectRef: project.ref, + orgSlug: selectedOrganization.slug, + reason, + spanId: state.messageSpanIds[messageId], + }) + + track('assistant_message_rating_submitted', { + rating, + category: result.category, + ...(reason && { reason }), + chatId, + }) + } catch (error) { + console.error('Failed to rate message:', error) + // Rollback on error + setMessageRatings((prev) => { + const { [messageId]: _, ...rest } = prev + return rest + }) + } + }, + [chatMessages, project?.ref, selectedOrganization?.slug, rateMessage, track, state, chatId] + ) + + const isContextExceededError = + error && + (error.message?.includes('context_length_exceeded') || + error.message?.includes('exceeds the context window')) + + const renderedMessages = useMemo( + () => + chatMessages.map((message, index) => { + const isBeingEdited = editingMessageId === message.id + const isAfterEditedMessage = editingMessageId + ? chatMessages.findIndex((m) => m.id === editingMessageId) < index + : false + const isLastMessage = index === chatMessages.length - 1 + + return ( + + + {branchedConversation && branchedFrom?.messageId === message.id && ( +
+
+
+ Branched from + +
+
+
+ )} + + ) + }), + [ + chatMessages, + deleteMessageFromHere, + editMessage, + cancelEdit, + editingMessageId, + chatStatus, + addToolApprovalResponse, + handleRateMessage, + messageRatings, + branchedConversation, + branchedFrom, + onSelectChat, + onBranchChat, + ] + ) + + const hasMessages = chatMessages.length > 0 + + const sendMessageToAssistant = (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) + } + + // Read off the attachments this message actually carries, so detaching the + // "Current Query" chip also drops the claim. Gated on the same condition that + // decides whether attachments make it into the text at all: with AI opt-in + // disabled the chip is shown but no query is sent, and claiming otherwise would + // have the server prepend ClickHouse context for a message holding no query. + // Rides on the message rather than the request, so a Retry reproduces the context + // the message was asked in. + const metadata: AssistantMessageMetadata = { + containsLogsSnippets: + includeSnippetsInMessage && containsLogsSnippets(composerContext?.sqlSnippets), + } + + const payload = { + role: 'user', + createdAt: new Date(), + parts: [{ type: 'text', text: finalContent }], + id: uuidv4(), + metadata, + } as MessageType + + composerContext?.onClearSqlSnippets?.() + lastUserMessageRef.current = payload + if (hasPendingApproval && !editingMessageId) { + setMessages(resolvePendingToolApprovalsAsDenied(chatMessages)) + } + sendMessage(payload, { + body: { + schema: currentSchema, + table: currentTable?.name, + }, + }) + setValue('') + + if (finalContent.includes('Help me to debug')) { + track('assistant_debug_submitted', { chatId }) + } else { + track('assistant_prompt_submitted', { chatId }) + } + } + + const handleClearMessages = () => { + if (isChatLoading) stop() + snap.clearMessages(chatId) + setMessages([]) + lastUserMessageRef.current = null + setEditingMessageId(null) + } + + useEffect(() => { + // Keep "Thinking" visible while stopping and resubmitting during edit + // Only clear once the new response actually starts streaming (or errors) + if (isResubmitting && (chatStatus === 'streaming' || !!error)) { + setIsResubmitting(false) + } + }, [isResubmitting, chatStatus, error]) + + useEffect(() => { + // Approval-required tools can't run in parallel. Auto-deny extras so the model reissues them sequentially. + for (const id of getParallelApprovalIdsToReject(chatMessages)) { + addToolApprovalResponse?.({ + id, + approved: false, + reason: + 'Only one approval-required tool call is allowed per turn. Please reissue this tool call after the current one completes.', + }) + } + }, [chatMessages, addToolApprovalResponse]) + + useEffect(() => { + setValue(composerContext?.initialInput || '') + if (inputRef.current && composerContext?.initialInput) { + inputRef.current.focus() + inputRef.current.setSelectionRange( + composerContext.initialInput.length, + composerContext.initialInput.length + ) + } + }, [composerContext?.initialInput]) + + return ( + { + handleClearMessages() + window.location.reload() + }, + }, + ]} + > +
+ {renderHeader?.({ + isChatLoading, + showMetadataWarning, + updatedOptInSinceMCP, + isHipaaProjectDisallowed, + aiOptInLevel, + })} + {hasMessages ? ( + + + {renderedMessages} + {error && ( + <> + + {isContextExceededError ? ( + + ) : ( + <> + + } + tooltip={{ content: { side: 'bottom', text: 'Clear messages' } }} + /> + + )} +
+ } + /> + + )} + {isChatLoading && ( + + )} + +

+ The Assistant can make mistakes. Double check responses. +

+ + + + ) : ( + setValue(val)} + onFocusInput={() => inputRef.current?.focus()} + /> + )} + + + {editingMessageId && ( + +
+ +
+
+ + Editing message +
+ } + onClick={cancelEdit} + className="w-6 h-6 p-0" + title="Cancel editing" + aria-label="Cancel editing" + tooltip={{ + content: { side: 'top', text: }, + }} + /> +
+ +
+ + )} + + +
+ {isSupportChat && !isSupportChatClosed && ( +
+
+
+ + +
+
+ )} + + {disablePrompts && ( + + )} + + {isSuccess && !isApiKeySet && ( + + } + /> + )} + + form>textarea]:text-base [&>form>textarea]:md:text-sm [&>form>textarea]:border', + '[&>form>textarea]:rounded-md [&>form>textarea]:outline-hidden!', + '[&>form>textarea]:ring-offset-0! [&>form>textarea]:ring-0!' + )} + loading={isChatLoading} + isEditing={!!editingMessageId} + disabled={isChatInputDisabled} + placeholder={ + hasMessages + ? isSupportChat + ? 'Share details so the assistant can help with your support request...' + : 'Ask a follow up question...' + : (composerContext?.sqlSnippets ?? []).length > 0 + ? 'Ask a question or make a change...' + : isSupportChat + ? 'Describe your support issue...' + : 'Chat to Postgres...' + } + value={value} + onValueChange={(e) => setValue(e.target.value)} + onSubmit={(finalMessage) => { + sendMessageToAssistant(finalMessage) + }} + onStop={() => { + stop() + // to save partial responses from the AI + const lastMessage = chatMessages[chatMessages.length - 1] + if (lastMessage && lastMessage.role === 'assistant') { + state.updateMessage(lastMessage, chatId) + } + }} + sqlSnippets={composerContext?.sqlSnippets} + onRemoveSnippet={(index) => { + const newSnippets = [...(composerContext?.sqlSnippets ?? [])] + newSnippets.splice(index, 1) + composerContext?.onSetSqlSnippets?.(newSnippets) + }} + includeSnippetsInMessage={includeSnippetsInMessage} + selectedModel={selectedModel} + onSelectModel={(model) => snap.setModel(model)} + /> +
+
+ + ) +} diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.Context.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.Context.tsx index 3ce0cb1dd57..d39824361ee 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.Context.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.Context.tsx @@ -26,6 +26,7 @@ export interface MessageActions { onDelete: (id: string) => void onEdit: (id: string) => void + onBranch: (id: string) => void onCancelEdit: () => void onRate?: (id: string, rating: 'positive' | 'negative', reason?: string) => void } diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx index c9a5ef5e6e5..6ec5780bec6 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.tsx @@ -8,11 +8,9 @@ import { MessageActions } from './Message.Actions' import type { AddToolApprovalResponse, MessageInfo } from './Message.Context' import { MessageProvider, useMessageActionsContext, useMessageInfoContext } from './Message.Context' import { MessageDisplay } from './Message.Display' -import { useAiAssistantStateSnapshot } from '@/state/ai-assistant-state' function AssistantMessage({ message }: { message: VercelMessage }) { - const snap = useAiAssistantStateSnapshot() - const { onCancelEdit, onRate } = useMessageActionsContext() + const { onBranch, onCancelEdit, onRate } = useMessageActionsContext() const { id, variant, state, isLastMessage, readOnly, rating, isLoading } = useMessageInfoContext() const handleRate = (newRating: 'positive' | 'negative', reason?: string) => { @@ -51,7 +49,7 @@ function AssistantMessage({ message }: { message: VercelMessage }) { isActive={rating === 'negative'} disabled={!!rating} /> - snap.branchChat(id)} /> + onBranch(id)} /> )} @@ -107,6 +105,7 @@ interface MessageProps { addToolApprovalResponse?: AddToolApprovalResponse onDelete: (id: string) => void onEdit: (id: string) => void + onBranch: (id: string) => void isAfterEditedMessage: boolean isBeingEdited: boolean onCancelEdit: () => void @@ -139,6 +138,7 @@ export function Message(props: MessageProps) { addToolApprovalResponse: props.addToolApprovalResponse, onDelete: props.onDelete, onEdit: props.onEdit, + onBranch: props.onBranch, onCancelEdit: props.onCancelEdit, onRate: props.onRate, } diff --git a/apps/studio/state/ai-assistant-state.test.ts b/apps/studio/state/ai-assistant-state.test.ts index 4fb412ffe83..d4670010878 100644 --- a/apps/studio/state/ai-assistant-state.test.ts +++ b/apps/studio/state/ai-assistant-state.test.ts @@ -1,7 +1,7 @@ import { proxy, ref } from 'valtio/vanilla' import { describe, expect, it } from 'vitest' -import { sanitizeForCloning } from './ai-assistant-state' +import { createAiAssistantState, sanitizeForCloning } from './ai-assistant-state' describe('AI assistant chat message sync', () => { // FE-3954: syncing the live array into valtio corrupted it with Proxies, breaking structuredClone in addToolApprovalResponse @@ -44,3 +44,82 @@ describe('AI assistant chat message sync', () => { expect(() => structuredClone(replacedMessage)).not.toThrow() }) }) + +describe('AI assistant chat surface isolation', () => { + it('creates chats without changing the sidebar selection', () => { + const state = createAiAssistantState() + const sidebarChatId = state.newChat({ name: 'Sidebar chat' }) + + const explorerChatId = state.createChat({ name: 'Explorer chat' }) + + expect(explorerChatId).not.toBe(sidebarChatId) + expect(state.activeChatId).toBe(sidebarChatId) + expect(state.chats[explorerChatId]?.name).toBe('Explorer chat') + }) + + it('branches a specified chat without changing the sidebar selection', () => { + const state = createAiAssistantState() + const sidebarChatId = state.newChat({ name: 'Sidebar chat' }) + const explorerChatId = state.createChat({ name: 'Explorer chat' }) + state.chats[explorerChatId].messages = [ + { id: 'message-1', role: 'user', parts: [{ type: 'text', text: 'Hello' }] }, + ] + + const branchId = state.createBranch(explorerChatId, 'message-1') + + expect(branchId).toBeDefined() + expect(state.activeChatId).toBe(sidebarChatId) + expect(state.chats[branchId!]?.branchedFrom).toEqual({ + chatId: explorerChatId, + messageId: 'message-1', + }) + }) + + it('mutates an explicit chat without changing or clearing the sidebar chat', () => { + const state = createAiAssistantState() + const sidebarChatId = state.newChat({ name: 'Sidebar chat' }) + const explorerChatId = state.createChat({ name: 'Explorer chat' }) + state.chats[sidebarChatId].messages = [ + { id: 'sidebar-message', role: 'user', parts: [{ type: 'text', text: 'Keep me' }] }, + ] + state.chats[explorerChatId].messages = [ + { id: 'explorer-message', role: 'user', parts: [{ type: 'text', text: 'Clear me' }] }, + ] + state.chatInstances[explorerChatId].messages = [ + { id: 'explorer-message', role: 'user', parts: [{ type: 'text', text: 'Clear me' }] }, + ] + + state.clearMessages(explorerChatId) + + expect(state.activeChatId).toBe(sidebarChatId) + expect(state.chats[sidebarChatId].messages).toHaveLength(1) + expect(state.chats[explorerChatId].messages).toHaveLength(0) + expect(state.chatInstances[explorerChatId].messages).toHaveLength(0) + }) + + it('keeps explicit chat message edits synchronized with the live chat instance', () => { + const state = createAiAssistantState() + const chatId = state.createChat({ name: 'Explorer chat' }) + const messages = [ + { id: 'message-1', role: 'user' as const, parts: [{ type: 'text' as const, text: 'First' }] }, + { + id: 'message-2', + role: 'user' as const, + parts: [{ type: 'text' as const, text: 'Second' }], + }, + ] + state.chats[chatId].messages = messages + state.chatInstances[chatId].messages = messages + + state.deleteMessagesAfter('message-2', { chatId }) + state.updateMessage( + { id: 'message-1', role: 'user', parts: [{ type: 'text', text: 'Updated' }] }, + chatId + ) + + expect(state.chats[chatId].messages).toEqual([ + { id: 'message-1', role: 'user', parts: [{ type: 'text', text: 'Updated' }] }, + ]) + expect(state.chatInstances[chatId].messages).toEqual(state.chats[chatId].messages) + }) +}) diff --git a/apps/studio/state/ai-assistant-state.tsx b/apps/studio/state/ai-assistant-state.tsx index 5355074cf65..8e0413bbc46 100644 --- a/apps/studio/state/ai-assistant-state.tsx +++ b/apps/studio/state/ai-assistant-state.tsx @@ -60,7 +60,7 @@ export type SupportChatMetadata = { isLifecycleSyncing: boolean } -type ChatSession = { +export type ChatSession = { id: string name: string messages: AssistantMessageType[] @@ -87,6 +87,10 @@ type AiAssistantData = { context: AiAssistantContext } +type CreateChatOptions = { name?: string; initialMessage?: string } +type NewChatOptions = CreateChatOptions & + Partial> + // Data structure stored in IndexedDB type StoredAiAssistantState = { projectRef: string @@ -389,11 +393,7 @@ export const createAiAssistantState = (): AiAssistantState => { return state.activeChatId ? state.chats[state.activeChatId] : undefined }, - newChat: ( - options?: { name?: string; initialMessage?: string } & Partial< - Pick - > - ) => { + createChat: (options?: CreateChatOptions) => { const chatId = uuidv4() const newChat: ChatSession = { id: chatId, @@ -407,21 +407,23 @@ export const createAiAssistantState = (): AiAssistantState => { ...state.chats, [chatId]: newChat, } - state.activeChatId = chatId - // Create new chat instance const chatInstance = createChatInstance(state, { id: chatId, initialMessages: [] }) - state.chatInstances[chatId] = ref(chatInstance) - // If initialMessage is provided, append it to the chat instance if (options?.initialMessage) { chatInstance.sendMessage({ text: options.initialMessage, }) } - // Update non-chat related state based on options, falling back to current state, then initial + return chatId + }, + + newChat: (options?: NewChatOptions) => { + const chatId = state.createChat(options) + state.selectChat(chatId) + const initialAiAssistantData = createInitialAiAssistantData() state.initialInput = options?.initialInput ?? initialAiAssistantData.initialInput state.sqlSnippets = options?.sqlSnippets ?? initialAiAssistantData.sqlSnippets @@ -431,8 +433,8 @@ export const createAiAssistantState = (): AiAssistantState => { return chatId }, - branchChat: (messageId: string) => { - const sourceChat = state.activeChat + createBranch: (sourceChatId: string, messageId: string) => { + const sourceChat = state.chats[sourceChatId] if (!sourceChat) return const messageIndex = sourceChat.messages.findIndex((msg) => msg.id === messageId) @@ -456,12 +458,22 @@ export const createAiAssistantState = (): AiAssistantState => { ...state.chats, [chatId]: newChat, } - state.activeChatId = chatId state.chatInstances[chatId] = ref( createChatInstance(state, { id: chatId, initialMessages: branchedMessages }) ) + return chatId + }, + + branchChat: (messageId: string) => { + if (!state.activeChatId) return + + const chatId = state.createBranch(state.activeChatId, messageId) + if (!chatId) return + + state.selectChat(chatId) + const initialAiAssistantData = createInitialAiAssistantData() state.initialInput = initialAiAssistantData.initialInput state.sqlSnippets = initialAiAssistantData.sqlSnippets @@ -493,35 +505,33 @@ export const createAiAssistantState = (): AiAssistantState => { }) }, - selectChat: (id: string) => { - if (id !== state.activeChatId) { - state.activeChatId = id - const chat = state.chats[id] - if (chat) { - if (!state.chatInstances[id]) { - state.chatInstances[id] = ref( - createChatInstance(state, { id, initialMessages: chat.messages }) - ) - } - } + ensureChatInstance: (id: string) => { + const chat = state.chats[id] + if (chat && !state.chatInstances[id]) { + state.chatInstances[id] = ref( + createChatInstance(state, { id, initialMessages: chat.messages }) + ) } }, + selectChat: (id: string) => { + if (!state.chats[id]) return + + state.activeChatId = id + state.ensureChatInstance(id) + }, + deleteChat: (id: string) => { const { [id]: _, ...remainingChats } = state.chats state.chats = remainingChats + delete state.chatInstances[id] if (id === state.activeChatId) { const remainingChatIds = Object.keys(remainingChats) state.activeChatId = remainingChatIds.length > 0 ? remainingChatIds[0] : undefined if (state.activeChatId) { - const chat = state.chats[state.activeChatId] - if (!state.chatInstances[state.activeChatId]) { - state.chatInstances[state.activeChatId] = ref( - createChatInstance(state, { id: state.activeChatId, initialMessages: chat.messages }) - ) - } + state.ensureChatInstance(state.activeChatId) } } }, @@ -534,19 +544,27 @@ export const createAiAssistantState = (): AiAssistantState => { } }, - clearMessages: () => { - const chat = state.activeChat + clearMessages: (chatId = state.activeChatId) => { + if (!chatId) return + + const chat = state.chats[chatId] if (chat) { chat.messages = [] + const chatInstance = state.chatInstances[chatId] + if (chatInstance) chatInstance.messages = [] chat.updatedAt = new Date() - state.suggestions = undefined - state.sqlSnippets = [] - state.initialInput = '' + if (chatId === state.activeChatId) { + state.suggestions = undefined + state.sqlSnippets = [] + state.initialInput = '' + } } }, - deleteMessagesAfter: (id: string, { includeSelf = true } = {}) => { - const chat = state.activeChat + deleteMessagesAfter: (id: string, { includeSelf = true, chatId = state.activeChatId } = {}) => { + if (!chatId) return + + const chat = state.chats[chatId] if (!chat) return const messageIndex = chat.messages.findIndex((msg) => msg.id === id) @@ -555,17 +573,37 @@ export const createAiAssistantState = (): AiAssistantState => { // Delete all messages from the target message (optionally including) to the end const startIndex = includeSelf ? messageIndex : messageIndex + 1 chat.messages.splice(startIndex) + const chatInstance = state.chatInstances[chatId] + const instanceMessageIndex = chatInstance?.messages.findIndex((message) => message.id === id) + if (chatInstance && instanceMessageIndex !== undefined && instanceMessageIndex !== -1) { + chatInstance.messages = chatInstance.messages.slice( + 0, + includeSelf ? instanceMessageIndex : instanceMessageIndex + 1 + ) + } chat.updatedAt = new Date() }, - updateMessage: (updatedMessage: MessageType) => { - const chat = state.activeChat + updateMessage: (updatedMessage: MessageType, chatId = state.activeChatId) => { + if (!chatId) return + + const chat = state.chats[chatId] if (!chat) return const messageIndex = chat.messages.findIndex((msg) => msg.id === updatedMessage.id) if (messageIndex !== -1) { // Clone first — valtio's proxy() mutates nested properties in place and would corrupt the SDK's live array - chat.messages[messageIndex] = sanitizeForCloning(updatedMessage) + const clonedMessage = sanitizeForCloning(updatedMessage) + chat.messages[messageIndex] = clonedMessage + const chatInstance = state.chatInstances[chatId] + const instanceMessageIndex = chatInstance?.messages.findIndex( + (message) => message.id === updatedMessage.id + ) + if (chatInstance && instanceMessageIndex !== undefined && instanceMessageIndex !== -1) { + chatInstance.messages = chatInstance.messages.map((message, index) => + index === instanceMessageIndex ? sanitizeForCloning(updatedMessage) : message + ) + } chat.updatedAt = new Date() } }, @@ -620,18 +658,7 @@ export const createAiAssistantState = (): AiAssistantState => { } // Initialize chat instance for the active chat - if ( - state.activeChatId && - state.chats[state.activeChatId] && - !state.chatInstances[state.activeChatId] - ) { - state.chatInstances[state.activeChatId] = ref( - createChatInstance(state, { - id: state.activeChatId, - initialMessages: state.chats[state.activeChatId].messages, - }) - ) - } + if (state.activeChatId) state.ensureChatInstance(state.activeChatId) }, clearStorage: async () => { @@ -650,19 +677,18 @@ export type AiAssistantState = AiAssistantData & { messageSpanIds: Record setContext: (context: Partial) => void setModel: (model: AssistantModel) => void - newChat: ( - options?: { name?: string; initialMessage?: string } & Partial< - Pick - > - ) => string + createChat: (options?: CreateChatOptions) => string + newChat: (options?: NewChatOptions) => string + createBranch: (sourceChatId: string, messageId: string) => string | undefined branchChat: (messageId: string) => string | undefined setSupportLifecycleStatus: (chatId: string, status: AiSupportStatus) => void + ensureChatInstance: (id: string) => void selectChat: (id: string) => void deleteChat: (id: string) => void renameChat: (id: string, name: string) => void - clearMessages: () => void - deleteMessagesAfter: (id: string, options?: { includeSelf?: boolean }) => void - updateMessage: (message: MessageType) => void + clearMessages: (chatId?: string) => void + deleteMessagesAfter: (id: string, options?: { includeSelf?: boolean; chatId?: string }) => void + updateMessage: (message: MessageType, chatId?: string) => void setSqlSnippets: (snippets: SqlSnippet[]) => void clearSqlSnippets: () => void loadPersistedState: (persistedState: StoredAiAssistantState) => void