mirror of
https://github.com/supabase/supabase.git
synced 2026-10-05 09:25:06 +03:00
refactor: decouple assistant chat surfaces (#48973)
## Summary - extract the reusable assistant conversation body from the sidebar interface - drive shared chat rendering and mutations with an explicit chat ID - keep sidebar selection, shortcuts, SQL-editor context, and header controls in the sidebar wrapper - add selection-free chat creation and branching primitives while preserving existing sidebar helpers - add store tests for cross-surface selection and mutation isolation This is PR 1 of 3. Review this foundation first, then #49031 for the routed Explorer chat lifecycle, followed by #49032 for discovery and cross-surface entry points. Supersedes #48970 after renaming the branch to use the repository's `chore/` convention. ## Test plan - `mise exec node@22 -- pnpm --dir apps/studio exec tsc --noEmit` - `mise exec node@22 -- pnpm --dir apps/studio exec vitest run state/ai-assistant-state.test.ts` - ESLint on changed files - Prettier on changed files
This commit is contained in:
Saxon Fletcher authored and GitHub committed
1 parent
00952806f2
commit
2f8181b8bb
8 files changed
+955
-732
No files matched your search
@@ -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<SqlSnippet, string>
|
||||
|
||||
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<AssistantModel>(() => {
|
||||
// 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<HTMLTextAreaElement>(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<MessageType | null>(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<string>(snap.initialInput || '')
|
||||
const [editingMessageId, setEditingMessageId] = useState<string | null>(null)
|
||||
const [isResubmitting, setIsResubmitting] = useState(false)
|
||||
const [messageRatings, setMessageRatings] = useState<Record<string, 'positive' | 'negative'>>({})
|
||||
|
||||
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 (
|
||||
<Fragment key={message.id}>
|
||||
<Message
|
||||
id={message.id}
|
||||
message={message}
|
||||
isLoading={chatStatus === 'submitted' || chatStatus === 'streaming'}
|
||||
readOnly={message.role === 'user'}
|
||||
addToolApprovalResponse={addToolApprovalResponse}
|
||||
onDelete={deleteMessageFromHere}
|
||||
onEdit={editMessage}
|
||||
isAfterEditedMessage={isAfterEditedMessage}
|
||||
isBeingEdited={isBeingEdited}
|
||||
onCancelEdit={cancelEdit}
|
||||
isLastMessage={isLastMessage}
|
||||
onRate={handleRateMessage}
|
||||
rating={messageRatings[message.id] ?? null}
|
||||
/>
|
||||
{branchedConversation && branchedFrom?.messageId === message.id && (
|
||||
<div className="flex items-center gap-2 mt-6">
|
||||
<div className="flex-1 border-t border-strong" />
|
||||
<div className="flex items-center gap-1 max-w-[80%] text-xs text-foreground-lighter">
|
||||
<span className="shrink-0">Branched from</span>
|
||||
<button
|
||||
tabIndex={0}
|
||||
className={cn(InlineLinkClassName, 'cursor-pointer truncate min-w-0')}
|
||||
onClick={() => snap.selectChat(branchedConversation.id)}
|
||||
>
|
||||
{branchedConversation.name}
|
||||
</button>
|
||||
</div>
|
||||
<div className="flex-1 border-t border-strong" />
|
||||
</div>
|
||||
)}
|
||||
</Fragment>
|
||||
)
|
||||
}),
|
||||
[
|
||||
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 (
|
||||
<ErrorBoundary
|
||||
message="Something went wrong with the AI Assistant"
|
||||
sentryContext={{
|
||||
component: 'AIAssistant',
|
||||
feature: 'AI Assistant Panel',
|
||||
projectRef: project?.ref,
|
||||
organizationSlug: selectedOrganization?.slug,
|
||||
<AssistantChat
|
||||
chatId={snap.activeChatId}
|
||||
className={className}
|
||||
shortcutsEnabled={shortcutsEnabled}
|
||||
onNewChat={handleNewChat}
|
||||
onSelectChat={(chatId) => 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()
|
||||
},
|
||||
},
|
||||
]}
|
||||
>
|
||||
<div className={cn('flex bg-card flex-col h-full w-full md:h-full max-h-dvh', className)}>
|
||||
renderHeader={(props) => (
|
||||
<AIAssistantHeader
|
||||
isChatLoading={isChatLoading}
|
||||
onNewChat={snap.newChat}
|
||||
{...props}
|
||||
shortcutsEnabled={shortcutsEnabled}
|
||||
onNewChat={handleNewChat}
|
||||
onCloseAssistant={() => closeSidebar(SIDEBAR_KEYS.AI_ASSISTANT)}
|
||||
showMetadataWarning={showMetadataWarning}
|
||||
updatedOptInSinceMCP={updatedOptInSinceMCP}
|
||||
isHipaaProjectDisallowed={isHipaaProjectDisallowed}
|
||||
aiOptInLevel={aiOptInLevel}
|
||||
/>
|
||||
{hasMessages ? (
|
||||
<Conversation className={cn('flex-1')}>
|
||||
<ConversationContent className="w-full px-7 py-8 mb-10 max-w-3xl mx-auto">
|
||||
{renderedMessages}
|
||||
{error && (
|
||||
<>
|
||||
<AlertError
|
||||
error={
|
||||
isContextExceededError
|
||||
? ASSISTANT_ERRORS['context-exceeded']
|
||||
: IS_PLATFORM
|
||||
? ASSISTANT_ERRORS['default']
|
||||
: error
|
||||
}
|
||||
showErrorPrefix={false}
|
||||
showInstructions={false}
|
||||
subject="Sorry, I'm having trouble responding right now."
|
||||
additionalActions={
|
||||
<div className="flex items-center gap-x-2 mr-auto">
|
||||
{isContextExceededError ? (
|
||||
<Button
|
||||
variant="default"
|
||||
size="tiny"
|
||||
onClick={() => snap.newChat()}
|
||||
className="text-xs"
|
||||
>
|
||||
New chat
|
||||
</Button>
|
||||
) : (
|
||||
<>
|
||||
<Button
|
||||
variant="default"
|
||||
size="tiny"
|
||||
onClick={() => regenerate()}
|
||||
className="text-xs"
|
||||
>
|
||||
Retry
|
||||
</Button>
|
||||
<ButtonTooltip
|
||||
variant="default"
|
||||
size="tiny"
|
||||
onClick={handleClearMessages}
|
||||
className="w-7 h-7"
|
||||
icon={<Eraser />}
|
||||
tooltip={{ content: { side: 'bottom', text: 'Clear messages' } }}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
{isChatLoading && (
|
||||
<motion.span
|
||||
animate={{ opacity: [1, 0] }}
|
||||
transition={{ duration: 1, repeat: Infinity, ease: 'linear' }}
|
||||
className="inline-block w-1.5 h-4 bg-foreground-lighter mt-4"
|
||||
/>
|
||||
)}
|
||||
|
||||
<p className="text-center text-xs text-foreground-muted mt-6">
|
||||
The Assistant can make mistakes. Double check responses.
|
||||
</p>
|
||||
</ConversationContent>
|
||||
<ConversationScrollButton />
|
||||
</Conversation>
|
||||
) : (
|
||||
<AIOnboarding
|
||||
key={snap.activeChatId}
|
||||
sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
|
||||
suggestions={
|
||||
snap.suggestions as
|
||||
| { title?: string; prompts?: { label: string; description: string }[] }
|
||||
| undefined
|
||||
}
|
||||
onValueChange={(val) => setValue(val)}
|
||||
onFocusInput={() => inputRef.current?.focus()}
|
||||
/>
|
||||
)}
|
||||
|
||||
<AnimatePresence>
|
||||
{editingMessageId && (
|
||||
<motion.div
|
||||
initial={{ opacity: 0 }}
|
||||
animate={{ opacity: 1 }}
|
||||
exit={{ opacity: 0 }}
|
||||
className="pointer-events-none z-10 -mt-24"
|
||||
>
|
||||
<div className="h-24 w-full bg-linear-to-t from-background to-transparent relative">
|
||||
<motion.div
|
||||
className="absolute left-1/2 z-20 bottom-8 pointer-events-auto"
|
||||
variants={{
|
||||
hidden: { y: 5, opacity: 0 },
|
||||
show: { y: 0, opacity: 1 },
|
||||
}}
|
||||
transition={{ duration: 0.1 }}
|
||||
initial="hidden"
|
||||
animate="show"
|
||||
exit="hidden"
|
||||
>
|
||||
<div className="-translate-x-1/2 bg-alternative dark:bg-muted border rounded-md px-3 py-2 min-w-[180px] flex items-center justify-between gap-x-2">
|
||||
<div className="flex items-center gap-x-2 text-sm text-foreground">
|
||||
<Pencil size={14} />
|
||||
<span>Editing message</span>
|
||||
</div>
|
||||
<ButtonTooltip
|
||||
variant="outline"
|
||||
size="tiny"
|
||||
icon={<X size={14} />}
|
||||
onClick={cancelEdit}
|
||||
className="w-6 h-6 p-0"
|
||||
title="Cancel editing"
|
||||
aria-label="Cancel editing"
|
||||
tooltip={{
|
||||
content: { side: 'top', text: <KeyboardShortcut keys={['Meta', 'Esc']} /> },
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
</motion.div>
|
||||
</div>
|
||||
</motion.div>
|
||||
)}
|
||||
</AnimatePresence>
|
||||
|
||||
<div className="px-3 pb-3 z-20 relative w-full max-w-3xl mx-auto flex flex-col gap-y-3">
|
||||
{isSupportChat && !isSupportChatClosed && (
|
||||
<div>
|
||||
<div className="mb-3 border-t" />
|
||||
<div className="flex items-center gap-2">
|
||||
<Button
|
||||
variant="outline"
|
||||
size="tiny"
|
||||
disabled={!activeChatId || !supportConversationId}
|
||||
onClick={() =>
|
||||
activeChatId && state.setSupportLifecycleStatus(activeChatId, 'escalated')
|
||||
}
|
||||
>
|
||||
Escalate to human
|
||||
</Button>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="tiny"
|
||||
disabled={!activeChatId || !supportConversationId}
|
||||
onClick={() =>
|
||||
activeChatId && state.setSupportLifecycleStatus(activeChatId, 'user_resolved')
|
||||
}
|
||||
>
|
||||
Resolve
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{disablePrompts && (
|
||||
<Admonition
|
||||
showIcon={false}
|
||||
type="default"
|
||||
title="Assistant has been temporarily disabled"
|
||||
description="We're currently looking into getting it back online"
|
||||
/>
|
||||
)}
|
||||
|
||||
{isSuccess && !isApiKeySet && (
|
||||
<Admonition
|
||||
type="default"
|
||||
title="OpenAI API key not set"
|
||||
description={
|
||||
<Markdown
|
||||
content={
|
||||
'Add your `OPENAI_API_KEY` to your environment variables to use the AI Assistant.'
|
||||
}
|
||||
/>
|
||||
}
|
||||
/>
|
||||
)}
|
||||
|
||||
<AssistantChatForm
|
||||
textAreaRef={inputRef}
|
||||
className={cn(
|
||||
'z-20',
|
||||
'[&>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)}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</ErrorBoundary>
|
||||
)}
|
||||
/>
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 = ({
|
||||
|
||||
<div className="flex items-center gap-x-4 shrink-0">
|
||||
<div className="flex items-center">
|
||||
<AIAssistantChatSelector />
|
||||
<AIAssistantChatSelector shortcutsEnabled={shortcutsEnabled} />
|
||||
|
||||
<ShortcutTooltip
|
||||
side="bottom"
|
||||
|
||||
@@ -0,0 +1,707 @@
|
||||
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 { Fragment, useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from 'react'
|
||||
import { Button, cn, KeyboardShortcut } from 'ui'
|
||||
import { Admonition } from 'ui-patterns/Admonition'
|
||||
|
||||
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 { 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 { 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 { SHORTCUT_IDS } from '@/state/shortcuts/registry'
|
||||
import { useShortcut } from '@/state/shortcuts/useShortcut'
|
||||
|
||||
export interface AssistantChatHeaderProps {
|
||||
isChatLoading: boolean
|
||||
showMetadataWarning: boolean
|
||||
updatedOptInSinceMCP: boolean
|
||||
isHipaaProjectDisallowed: boolean
|
||||
aiOptInLevel: 'disabled' | 'schema' | 'full' | string | undefined
|
||||
}
|
||||
|
||||
export interface AssistantChatComposerContext {
|
||||
initialInput?: string
|
||||
sqlSnippets?: SqlSnippet[]
|
||||
suggestions?: { title?: string; prompts?: { label: string; description: string }[] }
|
||||
onSetSqlSnippets?: (snippets: SqlSnippet[]) => 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<AssistantModel>(() => {
|
||||
// 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<HTMLTextAreaElement>(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<MessageType | null>(null)
|
||||
|
||||
const [value, setValue] = useState<string>(composerContext?.initialInput || '')
|
||||
const [editingMessageId, setEditingMessageId] = useState<string | null>(null)
|
||||
const [isResubmitting, setIsResubmitting] = useState(false)
|
||||
const [messageRatings, setMessageRatings] = useState<Record<string, 'positive' | 'negative'>>({})
|
||||
|
||||
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 (
|
||||
<Fragment key={message.id}>
|
||||
<Message
|
||||
id={message.id}
|
||||
message={message}
|
||||
isLoading={chatStatus === 'submitted' || chatStatus === 'streaming'}
|
||||
readOnly={message.role === 'user'}
|
||||
addToolApprovalResponse={addToolApprovalResponse}
|
||||
onDelete={deleteMessageFromHere}
|
||||
onEdit={editMessage}
|
||||
isAfterEditedMessage={isAfterEditedMessage}
|
||||
isBeingEdited={isBeingEdited}
|
||||
onCancelEdit={cancelEdit}
|
||||
isLastMessage={isLastMessage}
|
||||
onRate={handleRateMessage}
|
||||
rating={messageRatings[message.id] ?? null}
|
||||
onBranch={onBranchChat}
|
||||
/>
|
||||
{branchedConversation && branchedFrom?.messageId === message.id && (
|
||||
<div className="flex items-center gap-2 mt-6">
|
||||
<div className="flex-1 border-t border-strong" />
|
||||
<div className="flex items-center gap-1 max-w-[80%] text-xs text-foreground-lighter">
|
||||
<span className="shrink-0">Branched from</span>
|
||||
<button
|
||||
tabIndex={0}
|
||||
className={cn(InlineLinkClassName, 'cursor-pointer truncate min-w-0')}
|
||||
onClick={() => onSelectChat(branchedConversation.id)}
|
||||
>
|
||||
{branchedConversation.name}
|
||||
</button>
|
||||
</div>
|
||||
<div className="flex-1 border-t border-strong" />
|
||||
</div>
|
||||
)}
|
||||
</Fragment>
|
||||
)
|
||||
}),
|
||||
[
|
||||
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 (
|
||||
<ErrorBoundary
|
||||
message="Something went wrong with the AI Assistant"
|
||||
sentryContext={{
|
||||
component: 'AIAssistant',
|
||||
feature: 'AI Assistant Panel',
|
||||
projectRef: project?.ref,
|
||||
organizationSlug: selectedOrganization?.slug,
|
||||
}}
|
||||
actions={[
|
||||
{
|
||||
label: 'Clear messages and refresh',
|
||||
onClick: () => {
|
||||
handleClearMessages()
|
||||
window.location.reload()
|
||||
},
|
||||
},
|
||||
]}
|
||||
>
|
||||
<div className={cn('flex bg-card flex-col h-full w-full md:h-full max-h-dvh', className)}>
|
||||
{renderHeader?.({
|
||||
isChatLoading,
|
||||
showMetadataWarning,
|
||||
updatedOptInSinceMCP,
|
||||
isHipaaProjectDisallowed,
|
||||
aiOptInLevel,
|
||||
})}
|
||||
{hasMessages ? (
|
||||
<Conversation className={cn('flex-1')}>
|
||||
<ConversationContent className="w-full px-7 py-8 mb-10 max-w-3xl mx-auto">
|
||||
{renderedMessages}
|
||||
{error && (
|
||||
<>
|
||||
<AlertError
|
||||
error={
|
||||
isContextExceededError
|
||||
? ASSISTANT_ERRORS['context-exceeded']
|
||||
: IS_PLATFORM
|
||||
? ASSISTANT_ERRORS['default']
|
||||
: error
|
||||
}
|
||||
showErrorPrefix={false}
|
||||
showInstructions={false}
|
||||
subject="Sorry, I'm having trouble responding right now."
|
||||
additionalActions={
|
||||
<div className="flex items-center gap-x-2 mr-auto">
|
||||
{isContextExceededError ? (
|
||||
<Button
|
||||
variant="default"
|
||||
size="tiny"
|
||||
onClick={onNewChat}
|
||||
className="text-xs"
|
||||
>
|
||||
New chat
|
||||
</Button>
|
||||
) : (
|
||||
<>
|
||||
<Button
|
||||
variant="default"
|
||||
size="tiny"
|
||||
onClick={() => regenerate()}
|
||||
className="text-xs"
|
||||
>
|
||||
Retry
|
||||
</Button>
|
||||
<ButtonTooltip
|
||||
variant="default"
|
||||
size="tiny"
|
||||
onClick={handleClearMessages}
|
||||
className="w-7 h-7"
|
||||
icon={<Eraser />}
|
||||
tooltip={{ content: { side: 'bottom', text: 'Clear messages' } }}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
{isChatLoading && (
|
||||
<motion.span
|
||||
animate={{ opacity: [1, 0] }}
|
||||
transition={{ duration: 1, repeat: Infinity, ease: 'linear' }}
|
||||
className="inline-block w-1.5 h-4 bg-foreground-lighter mt-4"
|
||||
/>
|
||||
)}
|
||||
|
||||
<p className="text-center text-xs text-foreground-muted mt-6">
|
||||
The Assistant can make mistakes. Double check responses.
|
||||
</p>
|
||||
</ConversationContent>
|
||||
<ConversationScrollButton />
|
||||
</Conversation>
|
||||
) : (
|
||||
<AIOnboarding
|
||||
key={chatId}
|
||||
sqlSnippets={composerContext?.sqlSnippets}
|
||||
suggestions={composerContext?.suggestions}
|
||||
onValueChange={(val) => setValue(val)}
|
||||
onFocusInput={() => inputRef.current?.focus()}
|
||||
/>
|
||||
)}
|
||||
|
||||
<AnimatePresence>
|
||||
{editingMessageId && (
|
||||
<motion.div
|
||||
initial={{ opacity: 0 }}
|
||||
animate={{ opacity: 1 }}
|
||||
exit={{ opacity: 0 }}
|
||||
className="pointer-events-none z-10 -mt-24"
|
||||
>
|
||||
<div className="h-24 w-full bg-linear-to-t from-background to-transparent relative">
|
||||
<motion.div
|
||||
className="absolute left-1/2 z-20 bottom-8 pointer-events-auto"
|
||||
variants={{
|
||||
hidden: { y: 5, opacity: 0 },
|
||||
show: { y: 0, opacity: 1 },
|
||||
}}
|
||||
transition={{ duration: 0.1 }}
|
||||
initial="hidden"
|
||||
animate="show"
|
||||
exit="hidden"
|
||||
>
|
||||
<div className="-translate-x-1/2 bg-alternative dark:bg-muted border rounded-md px-3 py-2 min-w-[180px] flex items-center justify-between gap-x-2">
|
||||
<div className="flex items-center gap-x-2 text-sm text-foreground">
|
||||
<Pencil size={14} />
|
||||
<span>Editing message</span>
|
||||
</div>
|
||||
<ButtonTooltip
|
||||
variant="outline"
|
||||
size="tiny"
|
||||
icon={<X size={14} />}
|
||||
onClick={cancelEdit}
|
||||
className="w-6 h-6 p-0"
|
||||
title="Cancel editing"
|
||||
aria-label="Cancel editing"
|
||||
tooltip={{
|
||||
content: { side: 'top', text: <KeyboardShortcut keys={['Meta', 'Esc']} /> },
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
</motion.div>
|
||||
</div>
|
||||
</motion.div>
|
||||
)}
|
||||
</AnimatePresence>
|
||||
|
||||
<div className="px-3 pb-3 z-20 relative w-full max-w-3xl mx-auto flex flex-col gap-y-3">
|
||||
{isSupportChat && !isSupportChatClosed && (
|
||||
<div>
|
||||
<div className="mb-3 border-t" />
|
||||
<div className="flex items-center gap-2">
|
||||
<Button
|
||||
variant="outline"
|
||||
size="tiny"
|
||||
disabled={!supportConversationId}
|
||||
onClick={() => state.setSupportLifecycleStatus(chatId, 'escalated')}
|
||||
>
|
||||
Escalate to human
|
||||
</Button>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="tiny"
|
||||
disabled={!supportConversationId}
|
||||
onClick={() => state.setSupportLifecycleStatus(chatId, 'user_resolved')}
|
||||
>
|
||||
Resolve
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{disablePrompts && (
|
||||
<Admonition
|
||||
showIcon={false}
|
||||
type="default"
|
||||
title="Assistant has been temporarily disabled"
|
||||
description="We're currently looking into getting it back online"
|
||||
/>
|
||||
)}
|
||||
|
||||
{isSuccess && !isApiKeySet && (
|
||||
<Admonition
|
||||
type="default"
|
||||
title="OpenAI API key not set"
|
||||
description={
|
||||
<Markdown
|
||||
content={
|
||||
'Add your `OPENAI_API_KEY` to your environment variables to use the AI Assistant.'
|
||||
}
|
||||
/>
|
||||
}
|
||||
/>
|
||||
)}
|
||||
|
||||
<AssistantChatForm
|
||||
textAreaRef={inputRef}
|
||||
className={cn(
|
||||
'z-20',
|
||||
'[&>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)}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</ErrorBoundary>
|
||||
)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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}
|
||||
/>
|
||||
<MessageActions.Branch onClick={() => snap.branchChat(id)} />
|
||||
<MessageActions.Branch onClick={() => onBranch(id)} />
|
||||
</MessageActions>
|
||||
)}
|
||||
</MessageDisplay.Container>
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
@@ -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<Pick<AiAssistantData, 'initialInput' | 'sqlSnippets' | 'suggestions' | 'tables'>>
|
||||
|
||||
// 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<AiAssistantData, 'initialInput' | 'sqlSnippets' | 'suggestions' | 'tables'>
|
||||
>
|
||||
) => {
|
||||
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<string, string>
|
||||
setContext: (context: Partial<AiAssistantContext>) => void
|
||||
setModel: (model: AssistantModel) => void
|
||||
newChat: (
|
||||
options?: { name?: string; initialMessage?: string } & Partial<
|
||||
Pick<AiAssistantData, 'initialInput' | 'sqlSnippets' | 'suggestions' | 'tables'>
|
||||
>
|
||||
) => 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
|
||||
|
||||
Reference in new issue
Block a user