From 3fac43aac926712dae230ddd48a1b2d97b401d0e Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Wed, 17 Sep 2025 16:01:51 +1000 Subject: [PATCH] more refactoring --- .../Reports/ReportBlock/ReportBlock.tsx | 82 ++- .../ReportBlock/ReportBlockContainer.tsx | 16 +- .../ui/AIAssistantPanel/AIAssistant.tsx | 49 +- .../ui/AIAssistantPanel/AIOnboarding.tsx | 151 ++++- .../ui/AIAssistantPanel/ConfirmFooter.tsx | 6 +- .../AIAssistantPanel/DisplayBlockRenderer.tsx | 164 ++++-- .../ui/AIAssistantPanel/Message.tsx | 548 +++++++++--------- .../ui/AIAssistantPanel/MessageMarkdown.tsx | 122 +--- .../components/ui/QueryBlock/QueryBlock.tsx | 193 +++--- apps/studio/lib/ai/prompts.ts | 1 + apps/studio/pages/api/ai/sql/generate-v4.ts | 14 +- apps/studio/state/ai-assistant-state.tsx | 23 +- 12 files changed, 756 insertions(+), 613 deletions(-) diff --git a/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx b/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx index 642dfb65223..9675d768ccd 100644 --- a/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx +++ b/apps/studio/components/interfaces/Reports/ReportBlock/ReportBlock.tsx @@ -1,3 +1,4 @@ +import { useCallback, useEffect, useRef, useState } from 'react' import { X } from 'lucide-react' import { useParams } from 'common' @@ -6,13 +7,14 @@ import { ButtonTooltip } from 'components/ui/ButtonTooltip' import { DEFAULT_CHART_CONFIG, QueryBlock } from 'components/ui/QueryBlock/QueryBlock' import { AnalyticsInterval } from 'data/analytics/constants' import { useContentIdQuery } from 'data/content/content-id-query' +import { usePrimaryDatabase } from 'data/read-replicas/replicas-query' +import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' import { useDatabaseSelectorStateSnapshot } from 'state/database-selector' import { Dashboards, SqlSnippets } from 'types' -import { Button, cn } from 'ui' -import ShimmeringLoader from 'ui-patterns/ShimmeringLoader' import { DEPRECATED_REPORTS } from '../Reports.constants' import { ChartBlock } from './ChartBlock' import { DeprecatedChartBlock } from './DeprecatedChartBlock' +import { toast } from 'sonner' interface ReportBlockProps { item: Dashboards.Chart @@ -46,7 +48,7 @@ export const ReportBlock = ({ const isSnippet = item.attribute.startsWith('snippet_') - const { data, error, isLoading, isError } = useContentIdQuery( + const { data, error, isLoading } = useContentIdQuery( { projectRef, id: item.id }, { enabled: isSnippet && !!item.id, @@ -57,6 +59,11 @@ export const ReportBlock = ({ if (failureCount >= 2) return false return true }, + onSuccess: (contentData) => { + if (!isSnippet) return + const fetchedSql = (contentData?.content as SqlSnippets.Content | undefined)?.sql + if (fetchedSql) runQuery('select', fetchedSql) + }, } ) const sql = isSnippet ? (data?.content as SqlSnippets.Content)?.sql : undefined @@ -64,6 +71,66 @@ export const ReportBlock = ({ const isDeprecatedChart = DEPRECATED_REPORTS.includes(item.attribute) const snippetMissing = error?.message.includes('Content not found') + const { database: primaryDatabase } = usePrimaryDatabase({ projectRef }) + const readOnlyConnectionString = primaryDatabase?.connection_string_read_only + const postgresConnectionString = primaryDatabase?.connectionString + + const [rows, setRows] = useState(undefined) + const [isWriteQuery, setIsWriteQuery] = useState(false) + const lastQueryType = useRef<'select' | 'mutation'>('select') + + const { + mutate: executeSql, + error: executeSqlError, + isLoading: executeSqlLoading, + } = useExecuteSqlMutation({ + onSuccess: (data) => { + setRows(data.result) + setIsWriteQuery(lastQueryType.current === 'mutation') + }, + onError: (mutationError) => { + const lowerMessage = mutationError.message.toLowerCase() + const isReadOnlyError = + lowerMessage.includes('read-only transaction') || + lowerMessage.includes('permission denied') || + lowerMessage.includes('must be owner') + + if (lastQueryType.current === 'select' && isReadOnlyError) { + setIsWriteQuery(true) + } + }, + }) + + const runQuery = useCallback( + (queryType: 'select' | 'mutation' = 'select', sqlToRun?: string) => { + if (!projectRef || !sqlToRun) return false + + const connectionString = + queryType === 'mutation' + ? postgresConnectionString + : readOnlyConnectionString ?? postgresConnectionString + + if (!connectionString) { + toast.error('Unable to establish a database connection for this project.') + return false + } + + lastQueryType.current = queryType + if (queryType === 'mutation') { + setIsWriteQuery(true) + } + executeSql({ projectRef, connectionString, sql: sqlToRun }) + return true + }, + [projectRef, readOnlyConnectionString, postgresConnectionString, executeSql] + ) + + useEffect(() => { + if (!executeSqlLoading && isRefreshing) { + runQuery('select', sql) + } + }, [isRefreshing, executeSqlLoading, sql, runQuery]) + return ( <> {isSnippet ? ( @@ -72,6 +139,11 @@ export const ReportBlock = ({ label={item.label} chartConfig={chartConfig} sql={sql} + results={rows} + initialHideSql={true} + errorText={snippetMissing ? 'SQL snippet not found' : executeSqlError?.message} + isExecuting={executeSqlLoading} + isWriteQuery={isWriteQuery} actions={ !isLoading && ( ) } + onExecute={(queryType) => { + runQuery(queryType, sql) + }} onUpdateChartConfig={onUpdateChart} + disabled={isLoading || snippetMissing || !sql} /> ) : isDeprecatedChart ? ( )} -
- {children} +
+
+ {children} +
) diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index 40a5b7e0e48..babbde775da 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx @@ -106,6 +106,13 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { const { mutate: sendEvent } = useSendEventMutation() + const updateMessage = useCallback( + (updatedMessage: MessageType) => { + snap.updateMessage(updatedMessage) + }, + [snap] + ) + // Handle completion of the assistant's response const handleChatFinish = useCallback( ({ message }: { message: MessageType }) => { @@ -113,15 +120,10 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { snap.saveMessage([lastUserMessageRef.current, message]) lastUserMessageRef.current = null } else { - const existingMessage = snap.activeChat?.messages.find((m) => m.id === message.id) - if (existingMessage) { - snap.updateMessage(message) - } else { - snap.saveMessage(message) - } + updateMessage(message) } }, - [snap] + [snap, updateMessage] ) // TODO(refactor): This useChat hook should be moved down into each chat session. @@ -220,13 +222,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming' - const updateMessage = useCallback( - (updatedMessage: MessageType) => { - snap.updateMessage(updatedMessage) - }, - [snap] - ) - const deleteMessageFromHere = useCallback( (messageId: string) => { // Find the message index in current chatMessages @@ -422,7 +417,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { aiOptInLevel={aiOptInLevel} /> - {hasMessages && ( + {hasMessages ? ( {renderedMessages} {error && ( @@ -466,6 +461,17 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { /> )} + ) : ( + setValue(val)} + onFocusInput={() => inputRef.current?.focus()} + /> )} @@ -539,19 +545,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { /> )} - {isShowingOnboarding && ( - setValue(val)} - onFocusInput={() => inputRef.current?.focus()} - /> - )} - lint.level === LINTER_LEVELS.ERROR) ?? + []) as Lint[] + const securityErrorLints = errorLints.filter((lint) => lint.categories?.[0] === 'SECURITY') + const performanceErrorLints = errorLints.filter((lint) => lint.categories?.[0] !== 'SECURITY') + return ( -
-
-

How can I assist you?

-

- Generate SQL, RLS policies and edge functions, debug issues or check on your project - health. -

-
-
- {prompts.map((item, index) => ( - - - - ))} -
+ + + ))} +
+ ) : ( + <> + {performanceErrorLints.length > 0 && ( +
+

Improve Performance

+ {performanceErrorLints.map((lint, index) => { + return ( + + + + ) + })} +
+ )} + + {securityErrorLints.length > 0 && ( +
+

Improve Security

+ {securityErrorLints.map((lint, index) => { + return ( + + ) + })} +
+ )} + +
+

Ideas

+ {prompts.map((item, index) => ( + + + + ))} +
+ + )} ) } diff --git a/apps/studio/components/ui/AIAssistantPanel/ConfirmFooter.tsx b/apps/studio/components/ui/AIAssistantPanel/ConfirmFooter.tsx index 3567664372d..fdedd4548a5 100644 --- a/apps/studio/components/ui/AIAssistantPanel/ConfirmFooter.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/ConfirmFooter.tsx @@ -19,9 +19,9 @@ export const ConfirmFooter = ({ onConfirm, }: PropsWithChildren) => { return ( -
-
{message}
-
+
+
{message}
+
diff --git a/apps/studio/components/ui/AIAssistantPanel/DisplayBlockRenderer.tsx b/apps/studio/components/ui/AIAssistantPanel/DisplayBlockRenderer.tsx index b333042c266..7821d5ca4dd 100644 --- a/apps/studio/components/ui/AIAssistantPanel/DisplayBlockRenderer.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/DisplayBlockRenderer.tsx @@ -1,9 +1,11 @@ import { PermissionAction } from '@supabase/shared-types/out/constants' import { useRouter } from 'next/router' -import { DragEvent, PropsWithChildren, useEffect, useState } from 'react' +import { DragEvent, PropsWithChildren, useEffect, useRef, useState } from 'react' import { useParams } from 'common' import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartConfig' +import { usePrimaryDatabase } from 'data/read-replicas/replicas-query' +import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' import { useSendEventMutation } from 'data/telemetry/send-event-mutation' import { useAsyncCheckProjectPermissions } from 'hooks/misc/useCheckPermissions' import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization' @@ -12,8 +14,6 @@ import { useProfile } from 'lib/profile' import { ConfirmFooter } from './ConfirmFooter' import { DEFAULT_CHART_CONFIG, QueryBlock } from '../QueryBlock/QueryBlock' import { identifyQueryType } from './AIAssistant.utils' -import { usePrimaryDatabase } from 'data/read-replicas/replicas-query' -import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation' interface DisplayBlockRendererProps { messageId: string @@ -26,24 +26,30 @@ interface DisplayBlockRendererProps { xAxis?: string yAxis?: string } - results?: unknown - onResults: (args: { messageId: string; resultId?: string; results: unknown }) => void + initialResults?: unknown + onResults?: (args: { messageId: string; resultId?: string; results: unknown }) => void onError?: (args: { messageId: string; resultId?: string; errorText: string }) => void toolState?: 'input-streaming' | 'input-available' | 'output-available' | 'output-error' isLastPart?: boolean isLastMessage?: boolean + showConfirmFooter?: boolean + onChartConfigChange?: (chartConfig: ChartConfig) => void + onQueryRun?: (queryType: 'select' | 'mutation') => void } export const DisplayBlockRenderer = ({ messageId, toolCallId, initialArgs, - results, + initialResults, onResults, onError, toolState, isLastPart = false, isLastMessage = false, + showConfirmFooter = true, + onChartConfigChange, + onQueryRun, }: PropsWithChildren) => { const router = useRouter() const { ref } = useParams() @@ -68,29 +74,76 @@ export const DisplayBlockRenderer = ({ })) const resultId = toolCallId - const displayData = Array.isArray(results) ? results : undefined + const [rows, setRows] = useState( + Array.isArray(initialResults) ? initialResults : undefined + ) const isDraggableToReports = canCreateSQLSnippet && router.pathname.endsWith('/reports/[id]') const label = initialArgs.label || 'SQL Results' - const isWriteQuery = initialArgs.isWriteQuery || false + const [isWriteQuery, setIsWriteQuery] = useState(initialArgs.isWriteQuery || false) const sqlQuery = initialArgs.sql - 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({ + const lastQueryType = useRef<'select' | 'mutation'>('select') + + useEffect(() => { + const initialConfig: ChartConfig = { + ...DEFAULT_CHART_CONFIG, + view: initialArgs.view === 'chart' ? 'chart' : 'table', + xKey: initialArgs.xAxis ?? '', + yKey: initialArgs.yAxis ?? '', + } + setChartConfig(initialConfig) + onChartConfigChange?.(initialConfig) + setIsWriteQuery(initialArgs.isWriteQuery || false) + setRows(Array.isArray(initialResults) ? initialResults : undefined) + lastQueryType.current = initialArgs.isWriteQuery ? 'mutation' : 'select' + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [toolCallId]) + + useEffect(() => { + const normalized = Array.isArray(initialResults) ? initialResults : undefined + if (!normalized || normalized === rows) return + setRows(normalized) + }, [initialResults, rows]) + + const { + mutate: executeSql, + error: executeSqlError, + isLoading: executeSqlLoading, + } = useExecuteSqlMutation({ onSuccess: (data) => { - onResults({ messageId, resultId, results: data.result }) - setIsRunning(false) + setRows(Array.isArray(data.result) ? data.result : undefined) + setIsWriteQuery(lastQueryType.current === 'mutation' || initialArgs.isWriteQuery || false) + onResults?.({ + messageId, + resultId, + results: Array.isArray(data.result) ? data.result : undefined, + }) }, onError: (error) => { + const lowerMessage = error.message.toLowerCase() + const isReadOnlyError = + lowerMessage.includes('read-only transaction') || + lowerMessage.includes('permission denied') || + lowerMessage.includes('must be owner') + + if (lastQueryType.current === 'select' && isReadOnlyError) { + setIsWriteQuery(true) + } + onError?.({ messageId, resultId, errorText: error.message }) - setIsRunning(false) }, }) const handleRunQuery = (queryType: 'select' | 'mutation') => { + if (!sqlQuery) return + + onQueryRun?.(queryType) + sendEvent({ action: 'assistant_suggestion_run_query_clicked', properties: { @@ -104,12 +157,30 @@ export const DisplayBlockRenderer = ({ }) } + const runQuery = (queryType: 'select' | 'mutation') => { + if (!ref || !sqlQuery) return + + const connectionString = + queryType === 'mutation' + ? postgresConnectionString + : readOnlyConnectionString ?? postgresConnectionString + + if (!connectionString) { + const fallbackMessage = 'Unable to find a database connection to execute this query.' + onError?.({ messageId, resultId, errorText: fallbackMessage }) + return + } + + lastQueryType.current = queryType + if (queryType === 'mutation') { + setIsWriteQuery(true) + } + executeSql({ projectRef: ref, connectionString, sql: sqlQuery }) + } + const handleExecute = (queryType: 'select' | 'mutation') => { handleRunQuery(queryType) - if (ref && postgresConnectionString) { - setIsRunning(true) - execute({ projectRef: ref, connectionString: postgresConnectionString, sql: sqlQuery }) - } + runQuery(queryType) } const handleUpdateChartConfig = ({ @@ -117,7 +188,11 @@ export const DisplayBlockRenderer = ({ }: { chartConfig: Partial }) => { - setChartConfig((prev) => ({ ...prev, ...updatedValues })) + setChartConfig((prev) => { + const next = { ...prev, ...updatedValues } + onChartConfigChange?.(next) + return next + }) } const handleDragStart = (e: DragEvent) => { @@ -127,44 +202,41 @@ export const DisplayBlockRenderer = ({ ) } - const resolvedHasDecision = !!results + const resolvedHasDecision = initialResults !== undefined || rows !== undefined const shouldShowConfirmFooter = - !resolvedHasDecision && toolState === 'input-available' && isLastPart && isLastMessage + showConfirmFooter && + !resolvedHasDecision && + toolState === 'input-available' && + isLastPart && + isLastMessage return (
- { - onResults({ messageId, resultId, results }) - setIsRunning(false) - }} - onRunQuery={handleRunQuery} - onUpdateChartConfig={handleUpdateChartConfig} - onError={(errorText) => { - onError?.({ messageId, resultId, errorText }) - setIsRunning(false) - }} - draggable={isDraggableToReports} - onDragStart={handleDragStart} - disabled={!resolvedHasDecision || isRunning} - isExternallyExecuting={isRunning} - /> - +
+ +
{shouldShowConfirmFooter && (
{ - setIsRunning(false) - onResults({ messageId, resultId, results: 'User skipped running the query' }) + onResults?.({ messageId, resultId, results: 'User skipped running the query' }) }} onConfirm={() => { handleExecute(isWriteQuery ? 'mutation' : 'select') diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx index 9a5c239c470..18e57d6cdc6 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.tsx @@ -1,6 +1,6 @@ import { UIMessage as VercelMessage } from '@ai-sdk/react' import { CheckIcon, Loader2, Pencil, Trash2 } from 'lucide-react' -import { createContext, memo, PropsWithChildren, ReactNode, useMemo, useState } from 'react' +import { PropsWithChildren, ReactNode, useMemo, useState } from 'react' import ReactMarkdown from 'react-markdown' import { Components } from 'react-markdown/lib/ast-to-react' import remarkGfm from 'remark-gfm' @@ -23,13 +23,6 @@ import { OrderedList, } from './MessageMarkdown' import { Reasoning } from './elements/Reasoning' -import type { AssistantMessageType } from 'state/ai-assistant-state' - -interface MessageContextType { - isLoading: boolean - readOnly?: boolean -} -export const MessageContext = createContext({ isLoading: false }) const baseMarkdownComponents: Partial = { ol: OrderedList, @@ -45,7 +38,6 @@ interface MessageProps { message: VercelMessage isLoading: boolean readOnly?: boolean - action?: ReactNode variant?: 'default' | 'warning' onResults: (message: VercelMessage) => void addToolResult?: (args: { tool: string; toolCallId: string; output: unknown }) => Promise @@ -62,7 +54,6 @@ export const Message = function Message({ message, isLoading, readOnly, - action = null, variant = 'default', onResults, addToolResult, @@ -80,24 +71,12 @@ export const Message = function Message({ ...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, message] + [id, isLoading, readOnly] ) if (!message) { @@ -117,271 +96,302 @@ export const Message = function Message({ const hasTextContent = content && content.trim().length > 0 return ( - -
+
+ {isUser && ( + )} - onClick={isAfterEditedMessage ? onCancelEdit : undefined} - > - {variant === 'warning' && } - {action} +
+ {shouldUsePartsRendering ? ( + (() => { + 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' ? ( + + ) : ( + + )} +
+ {part.state === 'input-streaming' ? 'Running ' : 'Ran '} + {`${part.toolName}`} +
+
+ ) + } + case 'tool-list_policies': + case 'tool-search_docs': { + return ( +
+ {part.state === 'input-streaming' ? ( + + ) : ( + + )} +
+ {part.state === 'input-streaming' ? 'Running ' : 'Ran '} + {`${part.type.replace('tool-', '')}`} +
+
+ ) + } + case 'reasoning': + return ( + + {part.text || part.text} + + ) + case 'text': + return ( + 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' + )} + remarkPlugins={[remarkGfm]} + components={allMarkdownComponents} + > + {part.text} + + ) -
- {isUser && ( - - )} + case 'tool-execute_sql': { + const { toolCallId, state, input, output } = part + const inputArgs = input as any + const chartArgs = inputArgs?.chartConfig ?? inputArgs?.config ?? {} + const view = chartArgs?.view as 'table' | 'chart' | undefined + const xAxis = chartArgs?.xKey ?? chartArgs?.xAxis + const yAxis = chartArgs?.yKey ?? chartArgs?.yAxis -
- {shouldUsePartsRendering ? ( - (() => { - return parts.map( - (part: NonNullable[number], index: number) => { - const isLastPart = index === parts.length - 1 - switch (part.type) { - case 'dynamic-tool': { + if (state === 'input-streaming') { return (
- {part.state === 'input-streaming' ? ( - - ) : ( - - )} -
- {part.state === 'input-streaming' ? 'Running ' : 'Ran '} - {`${part.toolName}`} -
+ + Writing SQL...
) } - case 'reasoning': + + if (state === 'input-available' || state === 'output-available') { return ( - - {part.text || part.text} - - ) - case 'text': - return ( - 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' - )} - remarkPlugins={[remarkGfm]} - components={allMarkdownComponents} - > - {part.text} - - ) + { + const results = args.results as any[] - case 'tool-execute_sql': { - const { toolCallId, state, input, output } = part + console.log('onResults', args) - if (state === 'input-streaming') { - return ( -
- - Writing SQL... -
- ) - } - - if (state === 'input-available' || state === 'output-available') { - return ( -
- { - const results = args.results as any[] - - await addToolResult?.({ - tool: 'execute_sql', - toolCallId: String(toolCallId), - output: results, - }) - }} - onError={async ({ errorText }) => { - await addToolResult?.({ - tool: 'execute_sql', - toolCallId: String(toolCallId), - output: `Error: ${errorText}`, - }) - onResults(message) - }} - /> -
- ) - } - if (state === 'output-error') { - return ( -
- Failed to execute SQL. -
- ) - } - return null - } - case 'tool-deploy_edge_function': { - const { toolCallId, state, input } = part - if (state === 'input-streaming') { - return ( -
- - {`Preparing Edge Function deployment...`} -
- ) - } - if (state === 'input-available' || state === 'output-available') { - return ( - { - 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) - } + addToolResult?.({ + tool: 'execute_sql', + toolCallId: String(toolCallId), + output: results, + }) + }} + onError={({ errorText }) => { + addToolResult?.({ + tool: 'execute_sql', + toolCallId: String(toolCallId), + output: `Error: ${errorText}`, + }) + onResults(message) }} /> - ) - } - if (state === 'output-error') { - return ( -
- Failed to deploy Edge Function. -
- ) - } - return null +
+ ) } - - case 'source-url': - case 'source-document': - case 'file': - return null - default: - return null + if (state === 'output-error') { + return ( +
+ Failed to execute SQL. +
+ ) + } + return null + } + case 'tool-deploy_edge_function': { + const { toolCallId, state, input } = part + if (state === 'input-streaming') { + return ( +
+ + Writing Edge Function... +
+ ) + } + if (state === 'input-available' || state === 'output-available') { + return ( + { + 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') { + return ( +
+ Failed to deploy Edge Function. +
+ ) + } + return null } - } - ) - })() - ) : hasTextContent ? ( - - {content} - - ) : ( - Assistant is thinking... - )} -
-
- {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" - /> - -
-
- )} + case 'source-url': + case 'source-document': + case 'file': + return null + default: + return null + } + } + ) + })() + ) : hasTextContent ? ( + + {content} + + ) : ( + Assistant is thinking... + )} +
+ {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" + /> + +
+
+ )} { @@ -391,6 +401,6 @@ export const Message = function Message({ }} onCancel={() => setShowDeleteConfirmModal(false)} /> - +
) } diff --git a/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx b/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx index 5c781abfec3..5577eb8de60 100644 --- a/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx @@ -1,24 +1,11 @@ import { PermissionAction } from '@supabase/shared-types/out/constants' import { useRouter } from 'next/router' -import { - DragEvent, - memo, - ReactNode, - useCallback, - useContext, - useEffect, - useMemo, - useRef, -} from 'react' +import { memo, ReactNode, useEffect, useMemo, useRef } from 'react' import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartConfig' -import { useSendEventMutation } from 'data/telemetry/send-event-mutation' import { useAsyncCheckProjectPermissions } from 'hooks/misc/useCheckPermissions' -import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization' -import { useSelectedProjectQuery } from 'hooks/misc/useSelectedProject' import { useProfile } from 'lib/profile' import Link from 'next/link' -import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state' import { Button, cn, @@ -33,13 +20,10 @@ import { DialogTitle, DialogTrigger, } from 'ui' -import { DebouncedComponent } from '../DebouncedComponent' import { EdgeFunctionBlock } from '../EdgeFunctionBlock/EdgeFunctionBlock' -import { QueryBlock } from '../QueryBlock/QueryBlock' +import { DisplayBlockRenderer } from './DisplayBlockRenderer' import { AssistantSnippetProps } from './AIAssistant.types' -import { identifyQueryType } from './AIAssistant.utils' import { CollapsibleCodeBlock } from './CollapsibleCodeBlock' -import { MessageContext } from './Message' import { defaultUrlTransform } from './Message.utils' import { Loader2 } from 'lucide-react' @@ -126,27 +110,16 @@ Hyperlink.displayName = 'Hyperlink' export const MarkdownPre = ({ children, id, - onResults, + isLoading, + readOnly, }: { children: any id: string - onResults: ({ - messageId, - resultId, - results, - }: { - messageId: string - resultId?: string - results: any[] - }) => void + isLoading: boolean + readOnly?: boolean }) => { const router = useRouter() const { profile } = useProfile() - const { isLoading, readOnly } = useContext(MessageContext) - const { mutate: sendEvent } = useSendEventMutation() - const snap = useAiAssistantStateSnapshot() - const { data: project } = useSelectedProjectQuery() - const { data: org } = useSelectedOrganizationQuery() const { can: canCreateSQLSnippet } = useAsyncCheckProjectPermissions( PermissionAction.CREATE, @@ -183,12 +156,11 @@ export const MarkdownPre = ({ const snippetId = snippetProps.id const title = snippetProps.title || (language === 'edge' ? 'Edge Function' : 'SQL Query') const isChart = snippetProps.isChart === 'true' - const results = snap.getCachedSQLResults({ messageId: id, snippetId }) - // Strip props from the content for both SQL and edge functions const cleanContent = rawContent.replace(/(?:--|\/\/)\s*props:\s*\{[^}]+\}/, '').trim() const isDraggableToReports = canCreateSQLSnippet && router.pathname.endsWith('/reports/[id]') + const toolCallId = String(snippetId ?? id) useEffect(() => { chartConfig.current = { @@ -200,29 +172,6 @@ export const MarkdownPre = ({ // eslint-disable-next-line react-hooks/exhaustive-deps }, [snippetProps]) - const onResultsReturned = useCallback( - (results: any[]) => { - onResults({ messageId: id, resultId: snippetProps.id, results }) - }, - [onResults, snippetProps.id] - ) - - const onRunQuery = async (queryType: 'select' | 'mutation') => { - sendEvent({ - action: 'assistant_suggestion_run_query_clicked', - properties: { - queryType, - ...(queryType === 'mutation' - ? { category: identifyQueryType(cleanContent) ?? 'unknown' } - : {}), - }, - groups: { - project: project?.ref ?? 'Unknown', - organization: org?.slug ?? 'Unknown', - }, - }) - } - return (
{language === 'edge' ? ( @@ -235,42 +184,29 @@ export const MarkdownPre = ({ ) : language === 'sql' ? ( readOnly ? ( + ) : isLoading ? ( +
+ + Writing SQL... +
) : ( - - - Writing SQL... -
- } - > - { - chartConfig.current = { ...chartConfig.current, ...config } - }} - draggable={isDraggableToReports} - onDragStart={(e: DragEvent) => { - e.dataTransfer.setData( - 'application/json', - JSON.stringify({ label: title, sql: cleanContent, config: chartConfig.current }) - ) - }} - /> - + {}} + showConfirmFooter={false} + onChartConfigChange={(config) => { + chartConfig.current = { ...config } + }} + /> ) ) : ( void - onRunQuery?: (queryType: 'select' | 'mutation') => void - onError?: (errorText: string) => void + results?: any[] + errorText?: string + isExecuting?: boolean + initialHideSql?: boolean + onExecute?: (queryType: 'select' | 'mutation') => 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 @@ -56,74 +52,55 @@ export const QueryBlock = ({ sql, chartConfig = DEFAULT_CHART_CONFIG, actions, - initialResults, + results, + errorText, isWriteQuery = false, - onRunQuery, + isExecuting = false, + initialHideSql = false, + onExecute, onUpdateChartConfig, - onResults, - onError, draggable = false, onDragStart, disabled = false, - isExternallyExecuting = false, }: QueryBlockProps) => { - const { ref } = useParams() - const [chartSettings, setChartSettings] = useState(chartConfig) const { xKey, yKey, view = 'table' } = chartSettings - // 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 [showSql, setShowSql] = useState(!results && !initialHideSql) const [focusDataIndex, setFocusDataIndex] = useState() + const [showWarning, setShowWarning] = useState<'hasWriteOperation' | 'hasUnknownFunctions'>() + + const prevIsWriteQuery = useRef(isWriteQuery) + + useEffect(() => { + if (!prevIsWriteQuery.current && isWriteQuery) { + setShowWarning('hasWriteOperation') + } + if (!isWriteQuery && showWarning === 'hasWriteOperation') { + setShowWarning(undefined) + } + prevIsWriteQuery.current = isWriteQuery + }, [isWriteQuery, showWarning]) + + useEffect(() => { + setChartSettings(chartConfig) + }, [chartConfig]) const formattedQueryResult = useMemo(() => { - // Make sure Y axis values are numbers - return localResults?.map((row) => { + return results?.map((row) => { return Object.fromEntries( Object.entries(row).map(([key, value]) => { if (key === yKey) return [key, Number(value)] - else return [key, value] + return [key, value] }) ) }) - }, [localResults, yKey]) - - const [showWarning, setShowWarning] = useState<'hasWriteOperation' | 'hasUnknownFunctions'>() - - const { database: primaryDatabase } = usePrimaryDatabase({ projectRef: ref }) - const postgresConnectionString = primaryDatabase?.connectionString - const readOnlyConnectionString = primaryDatabase?.connection_string_read_only + }, [results, yKey]) const chartData = chartSettings.cumulative ? getCumulativeResults({ rows: formattedQueryResult ?? [] }, chartSettings) : formattedQueryResult - const { mutate: execute, isLoading: isInternallyExecuting } = useExecuteSqlMutation({ - onSuccess: (data) => { - setLocalResults(data.result) - onResults?.(data.result) - - setReadOnlyError(false) - setQueryError(undefined) - }, - onError: (error) => { - const readOnlyTransaction = /cannot execute .+ in a read-only transaction/.test(error.message) - const permissionDenied = error.message.includes('permission denied') - const notOwner = error.message.includes('must be owner') - if (readOnlyTransaction || permissionDenied || notOwner) { - setReadOnlyError(true) - setShowWarning('hasWriteOperation') - } else { - setQueryError(error) - } - onError?.(error.message) - }, - }) - const getDateFormat = (key: any) => { const value = chartData?.[0]?.[key] || '' if (typeof value === 'number') return 'number' @@ -132,39 +109,22 @@ export const QueryBlock = ({ } const xKeyDateFormat = getDateFormat(xKey) - // Combined loading state from internal execution or external execution - const isExecuting = isInternallyExecuting || isExternallyExecuting + const hasResults = Array.isArray(results) && results.length > 0 - const handleExecute = useCallback(() => { - if (!sql || isExecuting) return - - if (readOnlyError || isWriteQuery) { - return setShowWarning('hasWriteOperation') + const runSelect = () => { + if (!sql || disabled || isExecuting) return + if (isWriteQuery) { + setShowWarning('hasWriteOperation') + return } + onExecute?.('select') + } - try { - execute({ - projectRef: ref, - connectionString: readOnlyConnectionString, - sql, - }) - } catch (error: any) { - toast.error(`Failed to execute query: ${error.message}`) - } - }, [sql, isExecuting, readOnlyError, isWriteQuery, execute, ref, readOnlyConnectionString]) - - useEffect(() => { - setChartSettings(chartConfig) - }, [chartConfig]) - - // Sync local results when initialResults changes (from external execution) - useEffect(() => { - if (initialResults !== undefined) { - setLocalResults(initialResults) - } - }, [initialResults]) - - // Execution is only user-driven within this component + const runMutation = () => { + if (!sql || disabled || isExecuting) return + setShowWarning(undefined) + onExecute?.('mutation') + } return ( - {localResults && ( + {hasResults && ( { - if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: { view } }) - setChartSettings({ ...chartSettings, view }) + columns={Object.keys(results?.[0] ?? {})} + changeView={(nextView) => { + if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: { view: nextView } }) + setChartSettings({ ...chartSettings, view: nextView }) }} updateChartConfig={(config) => { if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: config }) @@ -212,20 +172,15 @@ export const QueryBlock = ({ className="w-7 h-7" icon={} loading={isExecuting} - disabled={isExecuting} - onClick={() => { - handleExecute() - if (!!sql) onRunQuery?.('select') - }} + disabled={isExecuting || disabled || !sql} + onClick={runSelect} tooltip={{ content: { side: 'bottom', className: 'max-w-56 text-center', - text: isExecuting ? ( -

{`Query is running. You may cancel ongoing queries via the [SQL Editor](/project/${ref}/sql?viewOngoingQueries=true).`}

- ) : ( - 'Run query' - ), + text: isExecuting + ? 'Query is running. Check the SQL Editor to manage running queries.' + : 'Run query', }, }} /> @@ -240,19 +195,7 @@ export const QueryBlock = ({ warningType={showWarning} className="border-b" onCancel={() => setShowWarning(undefined)} - onConfirm={() => { - // [Joshen] This is for when we introduced the concept of parameters into our reports - // const processedSql = processParameterizedSql(sql!, combinedParameterValues) - if (sql) { - setShowWarning(undefined) - execute({ - projectRef: ref, - connectionString: postgresConnectionString, - sql, - }) - onRunQuery?.('mutation') - } - }} + onConfirm={runMutation} disabled={!sql} {...(showWarning !== 'hasWriteOperation' ? { @@ -268,8 +211,8 @@ export const QueryBlock = ({ {showSql && (
)} - {isExecuting && localResults === undefined && ( + {isExecuting && !results && (
)} - {view === 'chart' && localResults !== undefined ? ( + {view === 'chart' && results !== undefined ? ( <> - {(localResults ?? []).length === 0 ? ( + {(results ?? []).length === 0 ? (

No results returned from query

@@ -305,7 +248,7 @@ export const QueryBlock = ({
) : ( <> - {!isExecuting && !!queryError ? ( + {!isExecuting && !!errorText ? (
- ERROR: {queryError.message} + ERROR: {errorText}
) : ( - localResults && ( + results && (
- +
) )} diff --git a/apps/studio/lib/ai/prompts.ts b/apps/studio/lib/ai/prompts.ts index c8c4bfb27ca..d391a0a3c1d 100644 --- a/apps/studio/lib/ai/prompts.ts +++ b/apps/studio/lib/ai/prompts.ts @@ -324,6 +324,7 @@ export const CHAT_PROMPT = ` # Project Health Checks - Use \`get_advisors\` to identify project issues. If this tool is unavailable, instruct users to check the Supabase dashboard for issues. +- Use \`get_logs\` to retrieve recent logs for the project # Safety for Destructive Queries - For destructive commands (e.g., DROP TABLE, DELETE without WHERE clause), always ask for confirmation before calling the \`execute_sql\` tool. diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index f4278684b96..a0168995d1c 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -109,6 +109,7 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { // 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) { + console.log('msg.parts', msg.parts) const cleanedParts = msg.parts .filter((part: any) => { // Filter out tool calls with invalid states for AI SDK @@ -119,9 +120,16 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { 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' } + // If opt in level is not schema_and_log_and_data, replace execute_sql tool output with permission message + if ( + aiOptInLevel !== 'schema_and_log_and_data' && + part.type === 'tool-execute_sql' && + part.output + ) { + return { + ...part, + output: 'The query was executed but the result is not shared due to permission level', + } } return part }) diff --git a/apps/studio/state/ai-assistant-state.tsx b/apps/studio/state/ai-assistant-state.tsx index 62e46842231..c97425903fc 100644 --- a/apps/studio/state/ai-assistant-state.tsx +++ b/apps/studio/state/ai-assistant-state.tsx @@ -339,15 +339,19 @@ export const createAiAssistantState = (): AiAssistantState => { const chat = state.activeChat if (!chat) return - const existingMessages = chat.messages - const messagesToAdd = Array.isArray(message) - ? message.filter( - (msg) => - !existingMessages.some((existing: AssistantMessageType) => existing.id === msg.id) - ) - : !existingMessages.some((existing: AssistantMessageType) => existing.id === message.id) - ? [message] - : [] + const incomingMessages = Array.isArray(message) ? message : [message] + + const messagesToAdd: AssistantMessageType[] = [] + + incomingMessages.forEach((msg) => { + const index = chat.messages.findIndex((existing) => existing.id === msg.id) + + if (index !== -1) { + state.updateMessage(msg) + } else { + messagesToAdd.push(msg as AssistantMessageType) + } + }) if (messagesToAdd.length > 0) { chat.messages.push(...messagesToAdd) @@ -361,7 +365,6 @@ export const createAiAssistantState = (): AiAssistantState => { const messageIndex = chat.messages.findIndex((msg) => msg.id === updatedMessage.id) if (messageIndex !== -1) { - // Replace entire message instance to capture latest tool parts/results chat.messages[messageIndex] = updatedMessage as AssistantMessageType chat.updatedAt = new Date() }