From 55f42c4eaeab79ca8854916c1f49cc4bb2bde0f2 Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Mon, 15 Sep 2025 14:57:57 +1000 Subject: [PATCH] all the ui changes --- .../Reports/ReportBlock/ReportBlock.tsx | 74 +--- .../ReportBlock/ReportBlockContainer.tsx | 31 +- .../ui/AIAssistantPanel/AIAssistant.tsx | 66 +-- .../ui/AIAssistantPanel/ConfirmFooter.tsx | 4 +- .../AIAssistantPanel/DisplayBlockRenderer.tsx | 130 +++--- .../AIAssistantPanel/EdgeFunctionRenderer.tsx | 79 ++++ .../ui/AIAssistantPanel/Message.tsx | 397 +++++------------- .../ui/AIAssistantPanel/MessageMarkdown.tsx | 24 +- .../components/ui/AIAssistantPanel/Tool.tsx | 0 .../components/ui/QueryBlock/QueryBlock.tsx | 300 ++++--------- apps/studio/lib/ai/prompts.ts | 28 +- apps/studio/lib/ai/tools/mcp-tools.ts | 7 +- apps/studio/lib/ai/tools/rendering-tools.ts | 8 +- apps/studio/pages/api/ai/sql/generate-v4.ts | 53 ++- apps/studio/state/ai-assistant-state.tsx | 45 +- 15 files changed, 503 insertions(+), 743 deletions(-) create mode 100644 apps/studio/components/ui/AIAssistantPanel/EdgeFunctionRenderer.tsx create mode 100644 apps/studio/components/ui/AIAssistantPanel/Tool.tsx diff --git a/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx b/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx index 103060a2bfa..642dfb65223 100644 --- a/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx +++ b/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx @@ -68,78 +68,22 @@ export const ReportBlock = ({ <> {isSnippet ? ( } - className="w-7 h-7" - onClick={() => onRemoveChart({ metric: { key: item.attribute } })} - tooltip={{ content: { side: 'bottom', text: 'Remove chart' } }} - /> + !isLoading && ( + } + className="w-7 h-7" + onClick={() => onRemoveChart({ metric: { key: item.attribute } })} + tooltip={{ content: { side: 'bottom', text: 'Remove chart' } }} + /> + ) } onUpdateChartConfig={onUpdateChart} - noResultPlaceholder={ -
- {isLoading ? ( - <> - - - - - ) : isError ? ( - <> -

- {snippetMissing ? 'SQL snippet cannot be found' : 'Error fetching SQL snippet'} -

-

- {snippetMissing ? 'Please remove this block from your report' : error.message} -

- - ) : ( - <> -

- No results returned from query -

-

- Results from the SQL query can be viewed as a table or chart here -

- - )} -
- } - readOnlyErrorPlaceholder={ -
-

- SQL query is not read-only and cannot be rendered -

-

- Queries that involve any mutation will not be run in reports -

- -
- } /> ) : isDeprecatedChart ? (
-
- {loading ? ( - - ) : ( - icon - )} -
- {showDragHandle && ( + {showDragHandle ? (
+ ) : ( + )} -

+

{label}

{actions}
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index facf502b3a8..40a5b7e0e48 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx @@ -37,7 +37,7 @@ import { ConversationContent, ConversationScrollButton, } from './elements/Conversation' -import { MemoizedMessage } from './Message' +import { Message } from './Message' interface AIAssistantProps { initialMessages?: MessageType[] | undefined @@ -113,7 +113,12 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { snap.saveMessage([lastUserMessageRef.current, message]) lastUserMessageRef.current = null } else { - snap.saveMessage(message) + const existingMessage = snap.activeChat?.messages.find((m) => m.id === message.id) + if (existingMessage) { + snap.updateMessage(message) + } else { + snap.saveMessage(message) + } } }, [snap] @@ -122,6 +127,17 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { // TODO(refactor): This useChat hook should be moved down into each chat session. // That way we won't have to disable switching chats while the chat is loading, // and don't run the risk of messages getting mixed up between chats. + // Sanitize messages to remove Valtio proxy wrappers that can't be cloned + const sanitizedMessages = useMemo(() => { + if (!snap.activeChat?.messages) return undefined + + return snap.activeChat.messages.map((msg: any) => { + // Convert proxy objects to plain objects + const plainMessage = JSON.parse(JSON.stringify(msg)) + return plainMessage + }) + }, [snap.activeChat?.messages]) + const { messages: chatMessages, status: chatStatus, @@ -133,26 +149,31 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { regenerate, } = useChat({ id: snap.activeChatId, - // [Alaister] typecast is needed here because valtio returns readonly arrays - // and useChat expects a mutable array sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls, - messages: snap.activeChat?.messages as unknown as MessageType[] | undefined, + messages: sanitizedMessages, async onToolCall({ toolCall }) { + if (toolCall.dynamic) { + return + } + if (toolCall.toolName === 'rename_chat') { const { newName } = toolCall.input as { newName: string } + if (snap.activeChatId && newName?.trim()) { snap.renameChat(snap.activeChatId, newName.trim()) + addToolResult({ tool: toolCall.toolName, toolCallId: toolCall.toolCallId, output: 'Chat renamed', }) + } else { + addToolResult({ + tool: toolCall.toolName, + toolCallId: toolCall.toolCallId, + output: 'Failed to rename chat: Invalid chat or name', + }) } - addToolResult({ - tool: toolCall.toolName, - toolCallId: toolCall.toolCallId, - output: 'Failed to rename chat: Invalid chat or name', - }) } }, transport: new DefaultChatTransport({ @@ -200,11 +221,10 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming' const updateMessage = useCallback( - (_args: { messageId: string; resultId?: string; results: any[] }) => { - // Intentionally no-op: we are not saving results into chat state - return + (updatedMessage: MessageType) => { + snap.updateMessage(updatedMessage) }, - [] + [snap] ) const deleteMessageFromHere = useCallback( @@ -266,25 +286,23 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { const isAfterEditedMessage = editingMessageId ? chatMessages.findIndex((m) => m.id === editingMessageId) < index : false + const isLastMessage = index === chatMessages.length - 1 return ( - Promise - } + addToolResult={addToolResult} onDelete={deleteMessageFromHere} onEdit={editMessage} isAfterEditedMessage={isAfterEditedMessage} isBeingEdited={isBeingEdited} onCancelEdit={cancelEdit} + isLastMessage={isLastMessage} /> ) }), @@ -444,7 +462,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { )} diff --git a/apps/studio/components/ui/AIAssistantPanel/ConfirmFooter.tsx b/apps/studio/components/ui/AIAssistantPanel/ConfirmFooter.tsx index 399b46111c3..3567664372d 100644 --- a/apps/studio/components/ui/AIAssistantPanel/ConfirmFooter.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/ConfirmFooter.tsx @@ -19,13 +19,13 @@ export const ConfirmFooter = ({ onConfirm, }: PropsWithChildren) => { return ( -
+
{message}
-
diff --git a/apps/studio/components/ui/AIAssistantPanel/DisplayBlockRenderer.tsx b/apps/studio/components/ui/AIAssistantPanel/DisplayBlockRenderer.tsx index 355b7995f22..b333042c266 100644 --- a/apps/studio/components/ui/AIAssistantPanel/DisplayBlockRenderer.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/DisplayBlockRenderer.tsx @@ -1,7 +1,6 @@ import { PermissionAction } from '@supabase/shared-types/out/constants' -import type { UIDataTypes, UIMessagePart, UITools } from 'ai' import { useRouter } from 'next/router' -import { DragEvent, PropsWithChildren, useMemo, useState } from 'react' +import { DragEvent, PropsWithChildren, useEffect, useState } from 'react' import { useParams } from 'common' import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartConfig' @@ -9,55 +8,47 @@ import { useSendEventMutation } from 'data/telemetry/send-event-mutation' import { useAsyncCheckProjectPermissions } from 'hooks/misc/useCheckPermissions' import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization' import { useProfile } from 'lib/profile' -import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state' -import { Badge } from 'ui' + +import { ConfirmFooter } from './ConfirmFooter' import { DEFAULT_CHART_CONFIG, QueryBlock } from '../QueryBlock/QueryBlock' import { identifyQueryType } from './AIAssistant.utils' -import { findResultForManualId } from './Message.utils' +import { usePrimaryDatabase } from 'data/read-replicas/replicas-query' +import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' interface DisplayBlockRendererProps { messageId: string toolCallId: string - manualId?: string initialArgs: { sql: string label?: string + isWriteQuery?: boolean view?: 'table' | 'chart' xAxis?: string yAxis?: string - runQuery?: boolean } - messageParts: UIMessagePart[] | undefined - isLoading: boolean - onResults: (args: { messageId: string; resultId?: string; results: any[] }) => void + results?: unknown + onResults: (args: { messageId: string; resultId?: string; results: unknown }) => void onError?: (args: { messageId: string; resultId?: string; errorText: string }) => void - /** - * External signal to trigger a query run. - * Increment this number to trigger another run. The component will - * translate this into a boolean toggle for QueryBlock.runQuery. - */ - triggerRunSignal?: number - /** Notify parent when a run starts or ends (for shared loading states) */ - onRunStateChange?: (isRunning: boolean) => void + toolState?: 'input-streaming' | 'input-available' | 'output-available' | 'output-error' + isLastPart?: boolean + isLastMessage?: boolean } export const DisplayBlockRenderer = ({ messageId, toolCallId, - manualId, initialArgs, - messageParts, - isLoading, + results, onResults, onError, - triggerRunSignal = 0, - onRunStateChange, + toolState, + isLastPart = false, + isLastMessage = false, }: PropsWithChildren) => { const router = useRouter() const { ref } = useParams() const { profile } = useProfile() const { data: org } = useSelectedOrganizationQuery() - const snap = useAiAssistantStateSnapshot() const { mutate: sendEvent } = useSendEventMutation() const { can: canCreateSQLSnippet } = useAsyncCheckProjectPermissions( @@ -76,21 +67,28 @@ export const DisplayBlockRenderer = ({ yKey: initialArgs.yAxis ?? '', })) - const isChart = initialArgs.view === 'chart' - const resultId = manualId || toolCallId - const liveResultData = useMemo( - () => (manualId ? findResultForManualId(messageParts, manualId) : undefined), - [messageParts, manualId] - ) - const cachedResults = useMemo( - () => snap.getCachedSQLResults({ messageId, snippetId: resultId }), - [snap, messageId, resultId] - ) - const displayData = liveResultData ?? cachedResults + const resultId = toolCallId + const displayData = Array.isArray(results) ? results : undefined const isDraggableToReports = canCreateSQLSnippet && router.pathname.endsWith('/reports/[id]') const label = initialArgs.label || 'SQL Results' + const isWriteQuery = initialArgs.isWriteQuery || false const sqlQuery = initialArgs.sql - const runQueryFlag = (triggerRunSignal || 0) % 2 === 1 + const [isRunning, setIsRunning] = useState(false) + + const { database: primaryDatabase } = usePrimaryDatabase({ projectRef: ref }) + const readOnlyConnectionString = primaryDatabase?.connection_string_read_only + const postgresConnectionString = primaryDatabase?.connectionString + + const { mutate: execute } = useExecuteSqlMutation({ + onSuccess: (data) => { + onResults({ messageId, resultId, results: data.result }) + setIsRunning(false) + }, + onError: (error) => { + onError?.({ messageId, resultId, errorText: error.message }) + setIsRunning(false) + }, + }) const handleRunQuery = (queryType: 'select' | 'mutation') => { sendEvent({ @@ -104,7 +102,14 @@ export const DisplayBlockRenderer = ({ organization: org?.slug ?? 'Unknown', }, }) - onRunStateChange?.(true) + } + + const handleExecute = (queryType: 'select' | 'mutation') => { + handleRunQuery(queryType) + if (ref && postgresConnectionString) { + setIsRunning(true) + execute({ projectRef: ref, connectionString: postgresConnectionString, sql: sqlQuery }) + } } const handleUpdateChartConfig = ({ @@ -122,42 +127,51 @@ export const DisplayBlockRenderer = ({ ) } + const resolvedHasDecision = !!results + const shouldShowConfirmFooter = + !resolvedHasDecision && toolState === 'input-available' && isLastPart && isLastMessage + return ( -
+
- - NEW - -

Drag to add this chart into your custom report

-
- ) : undefined - } onResults={(results) => { onResults({ messageId, resultId, results }) - onRunStateChange?.(false) + setIsRunning(false) }} onRunQuery={handleRunQuery} onUpdateChartConfig={handleUpdateChartConfig} onError={(errorText) => { onError?.({ messageId, resultId, errorText }) - onRunStateChange?.(false) + setIsRunning(false) }} + draggable={isDraggableToReports} onDragStart={handleDragStart} + disabled={!resolvedHasDecision || isRunning} + isExternallyExecuting={isRunning} /> + + {shouldShowConfirmFooter && ( +
+ { + setIsRunning(false) + onResults({ messageId, resultId, results: 'User skipped running the query' }) + }} + onConfirm={() => { + handleExecute(isWriteQuery ? 'mutation' : 'select') + }} + /> +
+ )}
) } diff --git a/apps/studio/components/ui/AIAssistantPanel/EdgeFunctionRenderer.tsx b/apps/studio/components/ui/AIAssistantPanel/EdgeFunctionRenderer.tsx new file mode 100644 index 00000000000..554db2ecdbc --- /dev/null +++ b/apps/studio/components/ui/AIAssistantPanel/EdgeFunctionRenderer.tsx @@ -0,0 +1,79 @@ +import { PropsWithChildren, useState } from 'react' +import { toast } from 'sonner' + +import { useParams } from 'common' +import { useEdgeFunctionDeployMutation } from 'data/edge-functions/edge-functions-deploy-mutation' +import { useSendEventMutation } from 'data/telemetry/send-event-mutation' +import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization' +import { EdgeFunctionBlock } from '../EdgeFunctionBlock/EdgeFunctionBlock' +import { ConfirmFooter } from './ConfirmFooter' + +interface EdgeFunctionRendererProps { + label: string + code: string + functionName: string + onDeployed?: (result: { success: true } | { success: false; errorText: string }) => void +} + +export const EdgeFunctionRenderer = ({ + label, + code, + functionName, + onDeployed, +}: PropsWithChildren) => { + const { ref } = useParams() + const { data: org } = useSelectedOrganizationQuery() + const { mutate: sendEvent } = useSendEventMutation() + const [isRunning, setIsRunning] = useState(false) + + const { mutateAsync: deployFunction } = useEdgeFunctionDeployMutation({ + onSuccess: () => { + toast.success('Successfully deployed edge function') + }, + }) + + const handleDeploy = async () => { + if (!ref || !functionName) return + try { + setIsRunning(true) + await deployFunction({ + projectRef: ref, + slug: functionName, + metadata: { + entrypoint_path: 'index.ts', + name: functionName, + verify_jwt: true, + }, + files: [{ name: 'index.ts', content: code }], + }) + sendEvent({ + action: 'edge_function_deploy_button_clicked', + properties: { origin: 'functions_ai_assistant' }, + groups: { project: ref ?? 'Unknown', organization: org?.slug ?? 'Unknown' }, + }) + onDeployed?.({ success: true }) + } catch (e: any) { + const errMsg = e?.message ?? 'Unknown error' + toast.error(`Failed to deploy function: ${errMsg}`) + onDeployed?.({ success: false, errorText: errMsg }) + } finally { + setIsRunning(false) + } + } + + return ( +
+ +
+ onDeployed?.({ success: false, errorText: 'Skipped' })} + onConfirm={handleDeploy} + /> +
+
+ ) +} diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx index 6c8580b8c6b..5a77f78ee36 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.tsx @@ -14,9 +14,11 @@ import { useProfile } from 'lib/profile' import { cn, markdownComponents, WarningIcon } from 'ui' import { ButtonTooltip } from '../ButtonTooltip' import { EdgeFunctionBlock } from '../EdgeFunctionBlock/EdgeFunctionBlock' +import { EdgeFunctionRenderer } from './EdgeFunctionRenderer' +import { ConfirmFooter } from './ConfirmFooter' import { DeleteMessageConfirmModal } from './DeleteMessageConfirmModal' import { DisplayBlockRenderer } from './DisplayBlockRenderer' -import { ConfirmFooter } from './ConfirmFooter' + import { Heading3, Hyperlink, @@ -26,6 +28,7 @@ import { OrderedList, } from './MessageMarkdown' import { Reasoning } from './elements/Reasoning' +import type { AssistantMessageType } from 'state/ai-assistant-state' interface MessageContextType { isLoading: boolean @@ -47,27 +50,19 @@ interface MessageProps { message: VercelMessage isLoading: boolean readOnly?: boolean - status?: string action?: ReactNode variant?: 'default' | 'warning' - onResults: ({ - messageId, - resultId, - results, - }: { - messageId: string - resultId?: string - results: any[] - }) => void + onResults: (message: VercelMessage) => void addToolResult?: (args: { tool: string; toolCallId: string; output: unknown }) => Promise onDelete: (id: string) => void onEdit: (id: string) => void isAfterEditedMessage: boolean isBeingEdited: boolean onCancelEdit: () => void + isLastMessage?: boolean } -const Message = function Message({ +export const Message = function Message({ id, message, isLoading, @@ -80,35 +75,35 @@ const Message = function Message({ onEdit, isAfterEditedMessage = false, isBeingEdited = false, - status, onCancelEdit, + isLastMessage = false, }: PropsWithChildren) { const { profile } = useProfile() - const { ref } = useParams() - const { data: project } = useSelectedProjectQuery() - const { mutateAsync: deployFunction } = useEdgeFunctionDeployMutation() const [showDeleteConfirmModal, setShowDeleteConfirmModal] = useState(false) - const [toolRunSignal, setToolRunSignal] = useState>({}) - const [toolRunning, setToolRunning] = useState>({}) const allMarkdownComponents: Partial = useMemo( () => ({ ...markdownComponents, ...baseMarkdownComponents, pre: ({ children }) => ( - + { + const updatedMessage: AssistantMessageType = { + ...(message as AssistantMessageType), + results: { + ...((message as AssistantMessageType).results ?? {}), + [args.resultId ?? 'default']: args.results as any[], + }, + } + onResults(updatedMessage) + }} + > {children} ), }), - [id, onResults] + [id, onResults, message] ) - const safeStringify = (value: unknown) => { - try { - return JSON.stringify(value) - } catch (_e) { - return 'null' - } - } if (!message) { console.error(`Message component received undefined message prop for id: ${id}`) @@ -130,7 +125,7 @@ const Message = function Message({
{ return parts.map( (part: NonNullable[number], index: number) => { + const isLastPart = index === parts.length - 1 switch (part.type) { case 'dynamic-tool': { return (
{part.state === 'input-streaming' ? ( @@ -175,7 +172,10 @@ const Message = function Message({ className="text-foreground-muted" /> )} - {`${part.toolName}`} +
+ {part.state === 'input-streaming' ? 'Running ' : 'Ran '} + {`${part.toolName}`} +
) } @@ -199,7 +199,7 @@ const Message = function Message({ div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h2:text-lg prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground prose-ol:space-y-3 prose-ul:space-y-3 prose-li:my-0 break-words [&>p:not(:last-child)]:!mb-2 [&>*>p:first-child]:!mt-0 [&>*>p:last-child]:!mb-0 [&>*>*>p:first-child]:!mt-0 [&>*>*>p:last-child]:!mb-0 [&>ol>li]:!pl-4', + 'max-w-none space-y-4 prose prose-sm prose-li:mt-1 [&>div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h2:text-lg prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground prose-ol:space-y-3 prose-ul:space-y-3 prose-li:my-0 break-words [&>p:not(:last-child)]:!mb-2 [&>*>p:first-child]:!mt-0 [&>*>p:last-child]:!mb-0 [&>*>*>p:first-child]:!mt-0 [&>*>*>p:last-child]:!mb-0 [&>ol>li]:!pl-4', isUser && 'text-foreground [&>p]:font-medium', isBeingEdited && 'animate-pulse' )} @@ -210,45 +210,17 @@ const Message = function Message({ ) - case 'tool-display_query': { - const { toolCallId, state, input } = part - if (state === 'input-streaming' || state === 'input-available') { - return ( -
- - {`Calling display_query...`} -
- ) - } - if (state === 'output-available') { - return ( - - ) - } - return null - } case 'tool-execute_sql': { - const { toolCallId, state, input } = part + const { toolCallId, state, input, output } = part + if (state === 'input-streaming') { return (
- {`Preparing SQL execution...`} + {'Writing: ' + (input as any)?.label || 'SQL Query'}
) } @@ -262,14 +234,27 @@ const Message = function Message({ { - addToolResult?.({ + initialArgs={{ + sql: (input as any).sql, + label: (input as any).label, + isWriteQuery: (input as any).isWriteQuery, + }} + results={output} + toolState={state} + isLastPart={isLastPart} + isLastMessage={isLastMessage} + onResults={async (args: { + messageId: string + resultId?: string + results: unknown + }) => { + const resultId = args.resultId + const results = args.results as any[] + + await addToolResult?.({ tool: 'execute_sql', toolCallId: String(toolCallId), - output: safeStringify(results), + output: results, }) }} onError={async ({ errorText }) => { @@ -278,38 +263,9 @@ const Message = function Message({ toolCallId: String(toolCallId), output: `Error: ${errorText}`, }) + onResults(message) }} - triggerRunSignal={toolRunSignal[toolCallId] || 0} - onRunStateChange={(isRunning) => - setToolRunning((prev) => ({ ...prev, [toolCallId]: isRunning })) - } /> - - {state === 'input-available' && ( -
- { - await addToolResult?.({ - tool: 'execute_sql', - toolCallId: String(toolCallId), - output: 'The user skipped running the query', - }) - setToolRunning((prev) => ({ ...prev, [toolCallId]: false })) - }} - onConfirm={() => { - setToolRunning((prev) => ({ ...prev, [toolCallId]: true })) - setToolRunSignal((prev) => ({ - ...prev, - [toolCallId]: (prev[toolCallId] || 0) + 1, - })) - }} - /> -
- )}
) } @@ -337,71 +293,29 @@ const Message = function Message({ } if (state === 'input-available' || state === 'output-available') { return ( -
- - {state === 'input-available' && ( -
- { - await addToolResult?.({ - tool: 'deploy_edge_function', - toolCallId: String(toolCallId), - output: 'Skipped', - }) - }} - onConfirm={async () => { - try { - setToolRunning((prev) => ({ - ...prev, - [toolCallId]: true, - })) - await deployFunction({ - projectRef: ref!, - slug: (input as any).name || 'my-function', - metadata: { - entrypoint_path: 'index.ts', - name: (input as any).name || 'my-function', - verify_jwt: true, - }, - files: [ - { name: 'index.ts', content: (input as any).code }, - ], - }) - await addToolResult?.({ - tool: 'deploy_edge_function', - toolCallId: String(toolCallId), - output: 'Deployed successfully', - }) - } catch (e: any) { - await addToolResult?.({ - tool: 'deploy_edge_function', - toolCallId: String(toolCallId), - output: `Error: ${e?.message ?? String(e)}`, - }) - } finally { - setToolRunning((prev) => ({ - ...prev, - [toolCallId]: false, - })) - } - }} - /> -
- )} -
+ label={(input as any).name || 'Edge Function'} + code={(input as any).code} + functionName={(input as any).name || 'my-function'} + onDeployed={async (res) => { + if (res.success) { + await addToolResult?.({ + tool: 'deploy_edge_function', + toolCallId: String(toolCallId), + output: 'Deployed successfully', + }) + onResults(message) + } else { + await addToolResult?.({ + tool: 'deploy_edge_function', + toolCallId: String(toolCallId), + output: res.errorText, + }) + onResults(message) + } + }} + /> ) } if (state === 'output-error') { @@ -413,36 +327,7 @@ const Message = function Message({ } return null } - case 'tool-display_edge_function': { - const { toolCallId, state, input } = part - if (state === 'input-streaming' || state === 'input-available') { - return ( -
- - {`Calling display_edge_function...`} -
- ) - } - if (state === 'output-available') { - return ( -
- -
- ) - } - return null - } + case 'source-url': case 'source-document': case 'file': @@ -464,44 +349,43 @@ const Message = function Message({ ) : ( Assistant is thinking... )} - - {/* Action button - only show for user messages on hover */} -
- {message.role === 'user' && ( - <> - } - onClick={ - isBeingEdited || isAfterEditedMessage ? onCancelEdit : () => onEdit(id) - } - className="text-foreground-light hover:text-foreground p-1 rounded" - aria-label={ - isBeingEdited || isAfterEditedMessage ? 'Cancel editing' : 'Edit message' - } - tooltip={{ - content: { - side: 'bottom', - text: - isBeingEdited || isAfterEditedMessage ? 'Cancel editing' : 'Edit message', - }, - }} - /> - - } - tooltip={{ content: { side: 'bottom', text: 'Delete message' } }} - onClick={() => setShowDeleteConfirmModal(true)} - className="text-foreground-light hover:text-foreground p-1 rounded" - title="Delete message" - aria-label="Delete message" - /> - - )} -
+ {message.role === 'user' && ( +
+ +
+ <> + } + onClick={isBeingEdited || isAfterEditedMessage ? onCancelEdit : () => onEdit(id)} + className="text-foreground-light hover:text-foreground p-1 rounded" + aria-label={ + isBeingEdited || isAfterEditedMessage ? 'Cancel editing' : 'Edit message' + } + tooltip={{ + content: { + side: 'bottom', + text: + isBeingEdited || isAfterEditedMessage ? 'Cancel editing' : 'Edit message', + }, + }} + /> + + } + tooltip={{ content: { side: 'bottom', text: 'Delete message' } }} + onClick={() => setShowDeleteConfirmModal(true)} + className="text-foreground-light hover:text-foreground p-1 rounded" + title="Delete message" + aria-label="Delete message" + /> + +
+
+ )}
) } - -export const MemoizedMessage = memo( - ({ - message, - status, - onResults, - addToolResult, - onDelete, - onEdit, - isAfterEditedMessage, - isBeingEdited, - onCancelEdit, - }: { - message: VercelMessage - status: string - onResults: ({ - messageId, - resultId, - results, - }: { - messageId: string - resultId?: string - results: any[] - }) => void - addToolResult?: (args: { tool: string; toolCallId: string; output: unknown }) => Promise - onDelete: (id: string) => void - onEdit: (id: string) => void - isAfterEditedMessage: boolean - isBeingEdited: boolean - onCancelEdit: () => void - }) => { - return ( - - ) - } -) - -MemoizedMessage.displayName = 'MemoizedMessage' diff --git a/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx b/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx index f3acd11c622..b3f0d92d766 100644 --- a/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx @@ -164,14 +164,12 @@ const MemoizedQueryBlock = memo( delay={isLoading ? 500 : 0} value={sql} fallback={ -
+
Writing SQL...
} > - - NEW - -

Drag to add this chart into your custom report

-
- ) : undefined - } - showSql={!isChart} - isChart={isChart} - isLoading={isLoading} - draggable={isDraggable} - runQuery={runQuery} - results={results} + initialResults={results} onRunQuery={onRunQuery} onResults={onResults} - onDragStart={onDragStart} onUpdateChartConfig={onUpdateChartConfig} + draggable={isDraggable} + onDragStart={onDragStart} /> ) diff --git a/apps/studio/components/ui/AIAssistantPanel/Tool.tsx b/apps/studio/components/ui/AIAssistantPanel/Tool.tsx new file mode 100644 index 00000000000..e69de29bb2d diff --git a/apps/studio/components/ui/QueryBlock/QueryBlock.tsx b/apps/studio/components/ui/QueryBlock/QueryBlock.tsx index 24ace51f948..bfa1a0454d3 100644 --- a/apps/studio/components/ui/QueryBlock/QueryBlock.tsx +++ b/apps/studio/components/ui/QueryBlock/QueryBlock.tsx @@ -10,16 +10,14 @@ import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartC import Results from 'components/interfaces/SQLEditor/UtilityPanel/Results' import { usePrimaryDatabase } from 'data/read-replicas/replicas-query' import { type QueryResponseError, useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' -import { type Parameter, parseParameters } from 'lib/sql-parameters' -import type { Dashboards } from 'types' -import { ChartContainer, ChartTooltipContent, cn, CodeBlock, SQL_ICON } from 'ui' + +import { Badge, ChartContainer, ChartTooltipContent, cn, CodeBlock, SQL_ICON } from 'ui' import ShimmeringLoader from 'ui-patterns/ShimmeringLoader' import { ButtonTooltip } from '../ButtonTooltip' import { CHART_COLORS } from '../Charts/Charts.constants' import SqlWarningAdmonition from '../SqlWarningAdmonition' import { BlockViewConfiguration } from './BlockViewConfiguration' import { EditQueryButton } from './EditQueryButton' -import { ParametersPopover } from './ParametersPopover' import { getCumulativeResults } from './QueryBlock.utils' export const DEFAULT_CHART_CONFIG: ChartConfig = { @@ -33,68 +31,21 @@ export const DEFAULT_CHART_CONFIG: ChartConfig = { } interface QueryBlockProps { - /** Applicable if SQL is a snippet that's already saved (Used in Reports) */ id?: string - /** Title of the QueryBlock */ label: string - /** SQL query to render/run in the QueryBlock */ sql?: string - /** Configuration of the output chart based on the query result */ + isWriteQuery?: boolean chartConfig?: ChartConfig - /** Not implemented yet: Will be the next part of ReportsV2 */ - parameterValues?: Record - /** Any other actions specific to the parent to be rendered in the header */ actions?: ReactNode - /** Toggle visiblity of SQL query on render */ - showSql?: boolean - /** Indicate if SQL query can be rendered as a chart */ - isChart?: boolean - /** For Assistant as QueryBlock is rendered while streaming response */ - isLoading?: boolean - /** Override to prevent running the SQL query provided */ - runQuery?: boolean - /** Prevent updating of columns for X and Y axes in the chart view */ - lockColumns?: boolean - /** Max height set to render results / charts (Defaults to 250) */ - maxHeight?: number - /** Whether query block is draggable */ - draggable?: boolean - /** Tooltip when hovering over the header of the block (Used in Assistant Panel) */ - tooltip?: ReactNode - /** Optional: Any initial results to render as part of the query*/ - results?: any[] - /** Opt to show run button if query is not read only */ - showRunButtonIfNotReadOnly?: boolean - /** Not implemented yet: Will be the next part of ReportsV2 */ - onSetParameter?: (params: Parameter[]) => void - /** Optional callback the SQL query is run */ - onRunQuery?: (queryType: 'select' | 'mutation') => void - /** Optional callback on drag start */ - onDragStart?: (e: DragEvent) => void - /** Optional: callback when the results are returned from running the SQL query*/ + initialResults?: any[] onResults?: (results: any[]) => void - /** Optional: callback when an error occurs while running the SQL query*/ + onRunQuery?: (queryType: 'select' | 'mutation') => void onError?: (errorText: string) => void - /** Force showing confirmation (write) bar before executing */ - forceShowConfirmation?: boolean - - // [Joshen] Params below are currently only used by ReportsV2 (Might revisit to see how to improve these) - /** Optional height set to render the SQL query (Used in Reports) */ - queryHeight?: number - /** UI to render if there's a read-only error while running the query */ - readOnlyErrorPlaceholder?: ReactNode - /** UI to render if there's no query results (Used in Reports) */ - noResultPlaceholder?: ReactNode - /** To trigger a refresh of the query */ - isRefreshing?: boolean - /** Optional callback whenever a chart configuration is updated (Used in Reports) */ - onUpdateChartConfig?: ({ - chart, - chartConfig, - }: { - chart?: Partial - chartConfig: Partial - }) => void + onUpdateChartConfig?: ({ chartConfig }: { chartConfig: Partial }) => void + draggable?: boolean + onDragStart?: (e: DragEvent) => void + disabled?: boolean + isExternallyExecuting?: boolean } // [Joshen ReportsV2] JFYI we may adjust this in subsequent PRs when we implement this into Reports V2 @@ -104,44 +55,34 @@ export const QueryBlock = ({ label, sql, chartConfig = DEFAULT_CHART_CONFIG, - maxHeight = 250, - queryHeight, - parameterValues: extParameterValues, actions, - showSql: _showSql = false, - isChart = false, - isLoading = false, - runQuery = false, - lockColumns = false, - draggable = false, - isRefreshing = false, - noResultPlaceholder = null, - readOnlyErrorPlaceholder = null, - showRunButtonIfNotReadOnly = false, - tooltip, - results, + initialResults, + isWriteQuery = false, onRunQuery, - onSetParameter, onUpdateChartConfig, - onDragStart, onResults, onError, - forceShowConfirmation, + draggable = false, + onDragStart, + disabled = false, + isExternallyExecuting = false, }: QueryBlockProps) => { const { ref } = useParams() const [chartSettings, setChartSettings] = useState(chartConfig) const { xKey, yKey, view = 'table' } = chartSettings - const [showSql, setShowSql] = useState(_showSql) + // Maintain local results state, seeded from initial results prop + const [localResults, setLocalResults] = useState(initialResults) + + const [showSql, setShowSql] = useState(!localResults) const [readOnlyError, setReadOnlyError] = useState(false) const [queryError, setQueryError] = useState() - const [queryResult, setQueryResult] = useState(results) const [focusDataIndex, setFocusDataIndex] = useState() const formattedQueryResult = useMemo(() => { // Make sure Y axis values are numbers - return queryResult?.map((row) => { + return localResults?.map((row) => { return Object.fromEntries( Object.entries(row).map(([key, value]) => { if (key === yKey) return [key, Number(value)] @@ -149,18 +90,10 @@ export const QueryBlock = ({ }) ) }) - }, [queryResult, yKey]) + }, [localResults, yKey]) - const [parameterValues, setParameterValues] = useState>({}) const [showWarning, setShowWarning] = useState<'hasWriteOperation' | 'hasUnknownFunctions'>() - const parameters = useMemo(() => { - if (!sql) return [] - return parseParameters(sql) - }, [sql]) - // [Joshen] This is for when we introduced the concept of parameters into our reports - // const combinedParameterValues = { ...extParameterValues, ...parameterValues } - const { database: primaryDatabase } = usePrimaryDatabase({ projectRef: ref }) const postgresConnectionString = primaryDatabase?.connectionString const readOnlyConnectionString = primaryDatabase?.connection_string_read_only @@ -169,10 +102,10 @@ export const QueryBlock = ({ ? getCumulativeResults({ rows: formattedQueryResult ?? [] }, chartSettings) : formattedQueryResult - const { mutate: execute, isLoading: isExecuting } = useExecuteSqlMutation({ + const { mutate: execute, isLoading: isInternallyExecuting } = useExecuteSqlMutation({ onSuccess: (data) => { + setLocalResults(data.result) onResults?.(data.result) - setQueryResult(data.result) setReadOnlyError(false) setQueryError(undefined) @@ -183,7 +116,7 @@ export const QueryBlock = ({ const notOwner = error.message.includes('must be owner') if (readOnlyTransaction || permissionDenied || notOwner) { setReadOnlyError(true) - if (showRunButtonIfNotReadOnly) setShowWarning('hasWriteOperation') + setShowWarning('hasWriteOperation') } else { setQueryError(error) } @@ -199,16 +132,17 @@ export const QueryBlock = ({ } const xKeyDateFormat = getDateFormat(xKey) + // Combined loading state from internal execution or external execution + const isExecuting = isInternallyExecuting || isExternallyExecuting + const handleExecute = useCallback(() => { if (!sql || isExecuting) return - if (readOnlyError) { + if (readOnlyError || isWriteQuery) { return setShowWarning('hasWriteOperation') } try { - // [Joshen] This is for when we introduced the concept of parameters into our reports - // const processedSql = processParameterizedSql(sql, combinedParameterValues) execute({ projectRef: ref, connectionString: readOnlyConnectionString, @@ -217,116 +151,72 @@ export const QueryBlock = ({ } catch (error: any) { toast.error(`Failed to execute query: ${error.message}`) } - }, [sql, isExecuting, readOnlyError, execute, ref, readOnlyConnectionString]) - - useEffect(() => { - if (forceShowConfirmation) setShowWarning('hasWriteOperation') - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [forceShowConfirmation]) + }, [sql, isExecuting, readOnlyError, isWriteQuery, execute, ref, readOnlyConnectionString]) useEffect(() => { setChartSettings(chartConfig) }, [chartConfig]) - // Run once on mount to parse parameters and notify parent + // Sync local results when initialResults changes (from external execution) useEffect(() => { - if (!!sql && onSetParameter) { - const params = parseParameters(sql) - onSetParameter(params) + if (initialResults !== undefined) { + setLocalResults(initialResults) } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [sql]) + }, [initialResults]) - useEffect(() => { - if (!!sql && runQuery && !!readOnlyConnectionString && !readOnlyError && !isExecuting) { - handleExecute() - } - }, [ - sql, - isLoading, - runQuery, - readOnlyConnectionString, - isExecuting, - readOnlyError, - handleExecute, - ]) - - useEffect(() => { - if (isRefreshing) handleExecute() - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [isRefreshing]) + // Execution is only user-driven within this component return ( ) => onDragStart?.(e)} - icon={ - + loading={isExecuting} + label={ +
+ {label} + {isWriteQuery && Write} +
} - label={label} actions={ - <> - } - onClick={() => setShowSql(!showSql)} - tooltip={{ - content: { side: 'bottom', text: showSql ? 'Hide query' : 'Show query' }, - }} - /> + disabled ? null : ( + <> + } + onClick={() => setShowSql(!showSql)} + tooltip={{ + content: { side: 'bottom', text: showSql ? 'Hide query' : 'Show query' }, + }} + /> + {localResults && ( + { + if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: { view } }) + setChartSettings({ ...chartSettings, view }) + }} + updateChartConfig={(config) => { + if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: config }) + setChartSettings(config) + }} + /> + )} - {queryResult && ( - <> - {/* [Joshen ReportsV2] Won't see this just yet as this is intended for Reports V2 */} - {parameters.length > 0 && ( - - )} - {isChart && ( - { - if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: { view } }) - setChartSettings({ ...chartSettings, view }) - }} - updateChartConfig={(config) => { - if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: config }) - setChartSettings(config) - }} - /> - )} - - )} - - - - {(showRunButtonIfNotReadOnly || !readOnlyError) && ( + } - loading={isExecuting || isLoading} - disabled={isLoading} + loading={isExecuting} + disabled={isExecuting} onClick={() => { handleExecute() if (!!sql) onRunQuery?.('select') @@ -343,10 +233,10 @@ export const QueryBlock = ({ }, }} /> - )} - {actions} - + {actions} + + ) } > {!!showWarning && ( @@ -383,9 +273,8 @@ export const QueryBlock = ({ {showSql && (
)} - {isExecuting && queryResult === undefined && ( + {isExecuting && localResults === undefined && (
)} - {view === 'chart' && queryResult !== undefined ? ( + {view === 'chart' && localResults !== undefined ? ( <> - {(queryResult ?? []).length === 0 ? ( + {(localResults ?? []).length === 0 ? (

No results returned from query

@@ -420,10 +309,7 @@ export const QueryBlock = ({
{!isExecuting && !!queryError ? ( -
+
ERROR: {queryError.message}
- ) : queryResult ? ( -
- -
- ) : !isExecuting ? ( - readOnlyError ? ( - readOnlyErrorPlaceholder - ) : ( - noResultPlaceholder + ) : ( + localResults && ( +
+ +
) - ) : null} + )} )} diff --git a/apps/studio/lib/ai/prompts.ts b/apps/studio/lib/ai/prompts.ts index 4b5e1d603c1..c8c4bfb27ca 100644 --- a/apps/studio/lib/ai/prompts.ts +++ b/apps/studio/lib/ai/prompts.ts @@ -280,41 +280,41 @@ Developer: # Role and Objective - Be aware that tool access may be restricted depending on the user's organization settings. - Do not try to bypass tool restrictions by executing SQL e.g. writing a query to retrieve database schema information. Instead, explain to the user you do not have permissions to use the tools you need to execute the task -# Output Format +## Output Format - Always integrate findings from the tools seamlessly into your responses for better accuracy and context. -# Searching Docs +## Searching Docs - Use \`search_docs\` to search the Supabase documentation for relevant information when the question is about Supabase features or complex database operations ` export const CHAT_PROMPT = ` -Developer: # Response Style -- Be direct and concise. Provide only essential information. -- Use lists to present information; do not use tables for formatting. -- Minimize use of emojis. +# Response Style +- Be professional, direct and concise. Provide only essential information. +- Before context gathering or tool usage, summarise the user's request and your plan of action in a single paragraph +- Do not repeat yourself or your plan after context gathering # Response Format -## Markdown -- Follow the CommonMark specification. -- Use a logical heading hierarchy (H1–H4), maintaining order without skipping levels. +## Use Markdown +- *CRITICAL*: Response must be in markdown format. +- Always use markdown blocks **where semantically correct** (e.g., \`inline code\`, \`\`\`code fences\`\`\`, headings, lists, tables). +- Make use of markdown headings to structure your response where appropriate. e.g. WRONG "Section heading: ..." WRITE "## Section heading" +- Shorter responses do not need headings - Use bold text exclusively to emphasize key information. - Do not use tables for displaying information under any circumstances. +- Minimize use of emojis. # Chat Naming - At the start of each conversation, if the chat has not yet been named, invoke \`rename_chat\` with a descriptive 2–4 word name. Examples: "User Authentication Setup", "Sales Data Analysis", "Product Table Creation". -## Task Workflow -- Always start the conversation with a concise checklist of sub-tasks you will perform before generating outputs or calling tools. Keep the checklist conceptual, not implementation-level. -- No need to repeat the checklist later in the conversation - # SQL Execution and Display - Be confident: assume the user is the project owner. You do not need to show code before execution. - To actually run SQL, directly call the \`execute_sql\` tool with the \`sql\` string. The client will request user confirmation and then return results. - If executing SQL returns an error, explain the error concisely and try again with the correct SQL. -- The user may skip executing the query in which case you should offer alternative options or actions to take +- The user may skip executing the query in which case you should acknowledge the skip and offer alternative options or actions to take - If the user asks you to write a query, or if you want to show example SQL without executing, render it in a markdown code block (e.g.: \`\`\`sql). Do this only when the user asks to see the code or for illustrative examples. - If multiple queries are needed, call \`execute_sql\` separately for each and validate results in 1–2 lines. Use separate code blocks only for non-executed examples. - After executing queries, summarize the outcome and confirm next actions or self-correct as needed. +- You do not need to repeat the SQL query results as the client will display them to the user as part of the execute_sql tool call. # Edge Functions - Be confident: assume the user is the project owner. You do not need to show code before deployment. diff --git a/apps/studio/lib/ai/tools/mcp-tools.ts b/apps/studio/lib/ai/tools/mcp-tools.ts index 777916f7693..eb136848744 100644 --- a/apps/studio/lib/ai/tools/mcp-tools.ts +++ b/apps/studio/lib/ai/tools/mcp-tools.ts @@ -3,6 +3,8 @@ import { createSupabaseMCPClient } from '../supabase-mcp' import { filterToolsByOptInLevel, toolSetValidationSchema } from '../tool-filter' import type { ToolSet } from 'ai' +const UI_EXECUTED_TOOLS = ['execute_sql', 'deploy_edge_function'] + export const getMcpTools = async ({ accessToken, projectRef, @@ -24,8 +26,9 @@ export const getMcpTools = async ({ // Remove UI-executed tools handled locally const filteredMcpTools: ToolSet = { ...allowedMcpTools } - delete (filteredMcpTools as any)['execute_sql'] - delete (filteredMcpTools as any)['deploy_edge_function'] + UI_EXECUTED_TOOLS.forEach((toolName) => { + delete filteredMcpTools[toolName] + }) // Validate that only known tools are provided const validation = toolSetValidationSchema.safeParse(filteredMcpTools) diff --git a/apps/studio/lib/ai/tools/rendering-tools.ts b/apps/studio/lib/ai/tools/rendering-tools.ts index 973f63f1b9a..96d06f631ef 100644 --- a/apps/studio/lib/ai/tools/rendering-tools.ts +++ b/apps/studio/lib/ai/tools/rendering-tools.ts @@ -6,8 +6,13 @@ export const getRenderingTools = () => ({ description: 'Asks the user to execute a SQL statement and return the results', inputSchema: z.object({ sql: z.string().describe('The SQL statement to execute.'), + label: z.string().describe('A short 2-4 word label for the SQL statement.'), + isWriteQuery: z + .boolean() + .describe( + 'Whether the SQL statement performs a write operation of any kind instead of a read operation' + ), }), - // No execute – runs client-side in the UI }), deploy_edge_function: tool({ description: @@ -16,7 +21,6 @@ export const getRenderingTools = () => ({ name: z.string().describe('The URL-friendly name/slug of the Edge Function.'), code: z.string().describe('The TypeScript code for the Edge Function.'), }), - // No execute – runs client-side in the UI }), rename_chat: tool({ description: `Rename the current chat session when the current chat name doesn't describe the conversation topic.`, diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index 91540460a04..f4278684b96 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -72,26 +72,6 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { const { messages: rawMessages, projectRef, connectionString, orgSlug, chatName } = data - // Server-side safety: limit to last 7 messages and remove `results` property to prevent accidental leakage. - // Results property is used to cache results client-side after queries are run - // Tool results will still be included in history sent to model - const messages = (rawMessages || []).slice(-7).map((msg: any) => { - if (msg && msg.role === 'assistant' && 'results' in msg) { - const cleanedMsg = { ...msg } - delete cleanedMsg.results - return cleanedMsg - } - // [Joshen] Am also filtering out any tool calls which state is "input-streaming" - // this happens when a user stops the assistant response while the tool is being called - if (msg && msg.role === 'assistant' && msg.parts) { - const cleanedParts = msg.parts.filter((part: any) => { - return !(part.type.startsWith('tool-') && part.state === 'input-streaming') - }) - return { ...msg, parts: cleanedParts } - } - return msg - }) - let aiOptInLevel: AiOptInLevel = 'disabled' let isLimited = false @@ -117,6 +97,39 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { } } + // Server-side safety: limit to last 7 messages and remove `results` property to prevent accidental leakage. + // Results property is used to cache results client-side after queries are run + // Tool results will still be included in history sent to model + const messages = (rawMessages || []).slice(-7).map((msg: any) => { + if (msg && msg.role === 'assistant' && 'results' in msg) { + const cleanedMsg = { ...msg } + delete cleanedMsg.results + return cleanedMsg + } + // Filter out tool calls with UI-specific states that the AI SDK doesn't understand + // These states are used for UI rendering but shouldn't be sent to the model + if (msg && msg.role === 'assistant' && msg.parts) { + const cleanedParts = msg.parts + .filter((part: any) => { + // Filter out tool calls with invalid states for AI SDK + if (part.type.startsWith('tool-')) { + const invalidStates = ['input-streaming', 'input-available', 'output-error'] + return !invalidStates.includes(part.state) + } + return true + }) + .map((part: any) => { + // If opt in level is schema_and_log_and_data, replace part output with permission message + if (aiOptInLevel !== 'schema_and_log_and_data' && part.output) { + return { ...part, output: 'Result not shared due to permission level' } + } + return part + }) + return { ...msg, parts: cleanedParts } + } + return msg + }) + const { model, error: modelError, diff --git a/apps/studio/state/ai-assistant-state.tsx b/apps/studio/state/ai-assistant-state.tsx index 4d2836c070c..62e46842231 100644 --- a/apps/studio/state/ai-assistant-state.tsx +++ b/apps/studio/state/ai-assistant-state.tsx @@ -104,16 +104,34 @@ async function clearStorage(): Promise { } } +// Helper function to sanitize objects to ensure they're cloneable +function sanitizeForCloning(obj: any): any { + if (obj === null || obj === undefined) return obj + if (typeof obj !== 'object') return obj + + // Convert to plain object to remove any proxy wrappers + return JSON.parse(JSON.stringify(obj)) +} + // Helper function to load state from IndexedDB async function loadFromIndexedDB(projectRef: string): Promise { try { const persistedState = await getAiState(projectRef) if (persistedState) { - // Revive dates + // Revive dates and sanitize message data Object.values(persistedState.chats).forEach((chat: ChatSession) => { if (chat && typeof chat === 'object') { chat.createdAt = new Date(chat.createdAt) chat.updatedAt = new Date(chat.updatedAt) + + // Sanitize message parts to remove proxy objects + if (chat.messages) { + chat.messages.forEach((message: any) => { + if (message.parts) { + message.parts = message.parts.map((part: any) => sanitizeForCloning(part)) + } + }) + } } }) return persistedState @@ -337,26 +355,15 @@ export const createAiAssistantState = (): AiAssistantState => { } }, - updateMessage: ({ - id, - resultId, - results, - }: { - id: string - resultId?: string - results: any[] - }) => { + updateMessage: (updatedMessage: MessageType) => { const chat = state.activeChat - if (!chat || !resultId) return - - const messageIndex = chat.messages.findIndex((msg) => msg.id === id) + if (!chat) return + const messageIndex = chat.messages.findIndex((msg) => msg.id === updatedMessage.id) if (messageIndex !== -1) { - const msg = chat.messages[messageIndex] - if (!msg.results) { - msg.results = {} - } - msg.results[resultId] = results + // Replace entire message instance to capture latest tool parts/results + chat.messages[messageIndex] = updatedMessage as AssistantMessageType + chat.updatedAt = new Date() } }, @@ -435,7 +442,7 @@ export type AiAssistantState = AiAssistantData & { clearMessages: () => void deleteMessagesAfter: (id: string, options?: { includeSelf?: boolean }) => void saveMessage: (message: MessageType | MessageType[]) => void - updateMessage: (args: { id: string; resultId?: string; results: any[] }) => void + updateMessage: (message: MessageType) => void setSqlSnippets: (snippets: SqlSnippet[]) => void clearSqlSnippets: () => void getCachedSQLResults: (args: { messageId: string; snippetId?: string }) => any[] | undefined