more refactoring

2026-10-13 13:19:09 +03:00 · 2025-09-17 16:01:51 +10:00
1 parent 1a91b561f6
commit 3fac43aac9
@@ -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<any[] | undefined>(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 && (
              <ButtonTooltip
@@ -83,7 +155,11 @@ export const ReportBlock = ({
              />
            )
          }
+          onExecute={(queryType) => {
+            runQuery(queryType, sql)
+          }}
          onUpdateChartConfig={onUpdateChart}
+          disabled={isLoading || snippetMissing || !sql}
        />
      ) : isDeprecatedChart ? (
        <DeprecatedChartBlock
@@ -65,8 +65,20 @@ export const ReportBlockContainer = ({
          </TooltipContent>
        )}
      </Tooltip>
-      <div className={cn('flex flex-col flex-grow items-center', hasChildren && 'border-t')}>
-        {children}
+      <div
+        className={cn(
+          'relative flex flex-col flex-grow w-full',
+          hasChildren && 'border-t overflow-hidden'
+        )}
+      >
+        <div
+          className={cn(
+            'flex flex-col flex-grow items-center overflow-hidden',
+            loading && 'pointer-events-none'
+          )}
+        >
+          {children}
+        </div>
      </div>
    </div>
  )
@@ -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 ? (
            <ConversationContent className="w-full px-7 py-8 mb-10">
              {renderedMessages}
              {error && (
@@ -466,6 +461,17 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
                />
              )}
            </ConversationContent>
+          ) : (
+            <AIOnboarding
+              sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
+              suggestions={
+                snap.suggestions as
+                  | { title?: string; prompts?: { label: string; description: string }[] }
+                  | undefined
+              }
+              onValueChange={(val) => setValue(val)}
+              onFocusInput={() => inputRef.current?.focus()}
+            />
          )}

          <ConversationScrollButton />
@@ -539,19 +545,6 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
            />
          )}

-          {isShowingOnboarding && (
-            <AIOnboarding
-              sqlSnippets={snap.sqlSnippets as SqlSnippet[] | undefined}
-              suggestions={
-                snap.suggestions as
-                  | { title?: string; prompts?: { label: string; description: string }[] }
-                  | undefined
-              }
-              onValueChange={(val) => setValue(val)}
-              onFocusInput={() => inputRef.current?.focus()}
-            />
-          )}
-
          <AssistantChatForm
            textAreaRef={inputRef}
            className={cn(
@@ -1,7 +1,11 @@
 import { motion } from 'framer-motion'
-import { FileText } from 'lucide-react'
+import { BarChart, FileText, Shield } from 'lucide-react'

 import { Button } from 'ui'
+import { useParams } from 'common'
+import { LINTER_LEVELS } from 'components/interfaces/Linter/Linter.constants'
+import { createLintSummaryPrompt, lintInfoMap } from 'components/interfaces/Linter/Linter.utils'
+import { type Lint, useProjectLintsQuery } from 'data/lint/lint-query'
 import { type SqlSnippet } from './AIAssistant.types'
 import { codeSnippetPrompts, defaultPrompts } from './AIAssistant.prompts'

@@ -31,38 +35,123 @@ export const AIOnboarding = ({
      ? codeSnippetPrompts
      : defaultPrompts

+  const { ref: projectRef } = useParams()
+  const { data: lints } = useProjectLintsQuery({ projectRef })
+
+  const errorLints: Lint[] = (lints?.filter((lint) => 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 (
-    <div className="w-full mb-6">
-      <div className="px-4 mb-4">
-        <h2 className="heading-section text-foreground mb-1">How can I assist you?</h2>
-        <p className="text-foreground-light text-sm">
-          Generate SQL, RLS policies and edge functions, debug issues or check on your project
-          health.
-        </p>
-      </div>
-      <div>
-        {prompts.map((item, index) => (
-          <motion.div
-            key={index}
-            initial={{ y: 5, opacity: 0 }}
-            animate={{ y: 0, opacity: 1 }}
-            transition={{ duration: 0.3, delay: index * 0.05 }}
-          >
-            <Button
-              size="small"
-              type="text"
-              className="w-full justify-start"
-              icon={<FileText strokeWidth={1.5} size={14} className="text-foreground-light" />}
-              onClick={() => {
-                onValueChange(item.prompt)
-                onFocusInput?.()
-              }}
+    <div className="w-full space-y-6 flex-1 overflow-auto py-8 px-4 justify-end flex flex-col gap-0">
+      <h2 className="heading-section text-foreground mx-4">How can I assist you?</h2>
+      {suggestions?.prompts?.length ? (
+        <div>
+          <h3 className="heading-meta text-foreground-light mb-3 mx-4">Suggestions</h3>
+          {prompts.map((item, index) => (
+            <motion.div
+              key={index}
+              initial={{ y: 5, opacity: 0 }}
+              animate={{ y: 0, opacity: 1 }}
+              transition={{ duration: 0.3, delay: index * 0.05 }}
            >
-              {item.title}
-            </Button>
-          </motion.div>
-        ))}
-      </div>
+              <Button
+                size="small"
+                type="text"
+                className="w-full justify-start border-b hover:border-b-0 hover:rounded-md rounded-none"
+                icon={<FileText strokeWidth={1.5} size={14} className="text-foreground-light" />}
+                onClick={() => {
+                  onValueChange(item.prompt)
+                  onFocusInput?.()
+                }}
+              >
+                {item.title}
+              </Button>
+            </motion.div>
+          ))}
+        </div>
+      ) : (
+        <>
+          {performanceErrorLints.length > 0 && (
+            <div className="mb-4">
+              <h3 className="heading-meta text-foreground-light mb-3 mx-4">Improve Performance</h3>
+              {performanceErrorLints.map((lint, index) => {
+                return (
+                  <motion.div
+                    key={lint.cache_key ?? index}
+                    initial={{ y: 5, opacity: 0 }}
+                    animate={{ y: 0, opacity: 1 }}
+                    transition={{ duration: 0.3, delay: index * 0.03 }}
+                  >
+                    <Button
+                      size="small"
+                      type="text"
+                      className="w-full justify-start"
+                      icon={
+                        <BarChart strokeWidth={1.5} size={14} className="text-foreground-light" />
+                      }
+                      onClick={() => {
+                        onValueChange(createLintSummaryPrompt(lint))
+                        onFocusInput?.()
+                      }}
+                    >
+                      {lint.detail ? lint.detail.replace('\\`', '') : lint.title}
+                    </Button>
+                  </motion.div>
+                )
+              })}
+            </div>
+          )}
+
+          {securityErrorLints.length > 0 && (
+            <div className="mb-4">
+              <h3 className="heading-meta text-foreground-light mb-3 mx-4">Improve Security</h3>
+              {securityErrorLints.map((lint, index) => {
+                return (
+                  <Button
+                    size="small"
+                    type="text"
+                    className="w-full justify-start"
+                    icon={<Shield strokeWidth={1.5} size={14} className="text-warning" />}
+                    onClick={() => {
+                      onValueChange(createLintSummaryPrompt(lint))
+                      onFocusInput?.()
+                    }}
+                  >
+                    {lint.detail ? lint.detail.replace(/`/g, '') : lint.title}
+                  </Button>
+                )
+              })}
+            </div>
+          )}
+
+          <div>
+            <h3 className="heading-meta text-foreground-light mb-3 mx-4">Ideas</h3>
+            {prompts.map((item, index) => (
+              <motion.div
+                key={index}
+                initial={{ y: 5, opacity: 0 }}
+                animate={{ y: 0, opacity: 1 }}
+                transition={{ duration: 0.3, delay: index * 0.05 }}
+              >
+                <Button
+                  size="small"
+                  type="text"
+                  className="w-full justify-start"
+                  icon={<FileText strokeWidth={1.5} size={14} className="text-foreground-light" />}
+                  onClick={() => {
+                    onValueChange(item.prompt)
+                    onFocusInput?.()
+                  }}
+                >
+                  {item.title}
+                </Button>
+              </motion.div>
+            ))}
+          </div>
+        </>
+      )}
    </div>
  )
 }
@@ -19,9 +19,9 @@ export const ConfirmFooter = ({
  onConfirm,
 }: PropsWithChildren<ConfirmFooterProps>) => {
  return (
-    <div className="rounded-b-lg border border-t-0 text-xs text-foreground p-2 pl-4 flex items-center justify-between gap-3  bg-gradient-to-r from-background-surface-75 to-background-muted shadow-inset">
-      <div className="flex-1">{message}</div>
-      <div className="flex items-center gap-2">
+    <div className="relative border border-t-0 overflow-hidden rounded-b-lg text-xs bg-border shadow-inset flex items-center justify-between gap-3 bg-gradient-to-r from-background-surface-75 to-background-surface-200 text-foreground py-2 pr-2 pl-4">
+      <div className="flex-1 relative z-10">{message}</div>
+      <div className="flex items-center gap-2 relative z-10">
        <Button size="tiny" type="outline" onClick={onCancel} disabled={isLoading}>
          {cancelLabel}
        </Button>
@@ -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<DisplayBlockRendererProps>) => {
  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<any[] | undefined>(
+    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<boolean>(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<ChartConfig>
  }) => {
-    setChartConfig((prev) => ({ ...prev, ...updatedValues }))
+    setChartConfig((prev) => {
+      const next = { ...prev, ...updatedValues }
+      onChartConfigChange?.(next)
+      return next
+    })
  }

  const handleDragStart = (e: DragEvent<Element>) => {
@@ -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 (
    <div className="display-block w-auto overflow-x-hidden">
-      <QueryBlock
-        label={label}
-        isWriteQuery={isWriteQuery}
-        sql={sqlQuery}
-        initialResults={displayData}
-        chartConfig={chartConfig}
-        onResults={(results) => {
-          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}
-      />
-
+      <div className="relative z-10">
+        <QueryBlock
+          label={label}
+          isWriteQuery={isWriteQuery}
+          sql={sqlQuery}
+          results={rows}
+          errorText={executeSqlError?.message}
+          chartConfig={chartConfig}
+          onExecute={handleExecute}
+          onUpdateChartConfig={handleUpdateChartConfig}
+          draggable={isDraggableToReports}
+          onDragStart={handleDragStart}
+          disabled={shouldShowConfirmFooter}
+          isExecuting={executeSqlLoading}
+        />
+      </div>
      {shouldShowConfirmFooter && (
        <div className="mx-4">
          <ConfirmFooter
            message="Assistant wants to run this query"
            cancelLabel="Skip"
-            confirmLabel={isRunning ? 'Running…' : 'Run Query'}
-            isLoading={isRunning}
+            confirmLabel={executeSqlLoading ? 'Running…' : 'Run Query'}
+            isLoading={executeSqlLoading}
            onCancel={async () => {
-              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')
@@ -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<MessageContextType>({ isLoading: false })

 const baseMarkdownComponents: Partial<Components> = {
  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<void>
@@ -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 }) => (
-        <MarkdownPre
-          id={id}
-          onResults={(args: { messageId: string; resultId?: string; results: any[] }) => {
-            const updatedMessage: AssistantMessageType = {
-              ...(message as AssistantMessageType),
-              results: {
-                ...((message as AssistantMessageType).results ?? {}),
-                [args.resultId ?? 'default']: args.results as any[],
-              },
-            }
-            onResults(updatedMessage)
-          }}
-        >
+        <MarkdownPre id={id} isLoading={isLoading} readOnly={readOnly}>
          {children}
        </MarkdownPre>
      ),
    }),
-    [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 (
-    <MessageContext.Provider value={{ isLoading, readOnly }}>
-      <div
-        className={cn(
-          'group text-foreground-light text-sm first:mt-0',
-          isUser ? 'text-foreground mt-6' : '',
-          variant === 'warning' && 'bg-warning-200',
-          isAfterEditedMessage && 'opacity-50 cursor-pointer transition-opacity'
+    <div
+      className={cn(
+        'group text-foreground-light text-sm first:mt-0',
+        isUser ? 'text-foreground mt-6' : '',
+        variant === 'warning' && 'bg-warning-200',
+        isAfterEditedMessage && 'opacity-50 cursor-pointer transition-opacity'
+      )}
+      onClick={isAfterEditedMessage ? onCancelEdit : undefined}
+    >
+      <div className="flex gap-4 w-auto overflow-hidden group">
+        {isUser && (
+          <ProfileImage
+            alt={profile?.username}
+            src={profile?.profileImageUrl}
+            className="w-5 h-5 shrink-0 rounded-full translate-y-0.5"
+          />
        )}
-        onClick={isAfterEditedMessage ? onCancelEdit : undefined}
-      >
-        {variant === 'warning' && <WarningIcon className="w-6 h-6" />}

-        {action}
+        <div className="flex-1 min-w-0">
+          {shouldUsePartsRendering ? (
+            (() => {
+              return parts.map(
+                (part: NonNullable<VercelMessage['parts']>[number], index: number) => {
+                  const isLastPart = index === parts.length - 1
+                  switch (part.type) {
+                    case 'dynamic-tool': {
+                      return (
+                        <div
+                          key={`${id}-tool-${part.toolCallId}`}
+                          className={cn(
+                            'tool-item text-foreground-lighter flex items-center gap-2 py-2',
+                            '[&:not(.tool-item+.tool-item)]:mt-4 [&:not(:has(+.tool-item))]:mb-4',
+                            '[&:has(+.tool-item)]:border-b [&:has(+.tool-item)]:border-b-muted',
+                            'first:!mt-0 last:mb-0'
+                          )}
+                        >
+                          {part.state === 'input-streaming' ? (
+                            <Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
+                          ) : (
+                            <CheckIcon
+                              strokeWidth={1.5}
+                              size={12}
+                              className="text-foreground-muted"
+                            />
+                          )}
+                          <div>
+                            {part.state === 'input-streaming' ? 'Running ' : 'Ran '}
+                            <span className="text-foreground-lighter">{`${part.toolName}`}</span>
+                          </div>
+                        </div>
+                      )
+                    }
+                    case 'tool-list_policies':
+                    case 'tool-search_docs': {
+                      return (
+                        <div
+                          key={`${id}-tool-${part.toolCallId}`}
+                          className={cn(
+                            'tool-item text-foreground-lighter flex items-center gap-2 py-2',
+                            '[&:not(.tool-item+.tool-item)]:mt-4 [&:not(:has(+.tool-item))]:mb-4',
+                            '[&:has(+.tool-item)]:border-b [&:has(+.tool-item)]:border-b-muted',
+                            'first:!mt-0 last:mb-0'
+                          )}
+                        >
+                          {part.state === 'input-streaming' ? (
+                            <Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
+                          ) : (
+                            <CheckIcon
+                              strokeWidth={1.5}
+                              size={12}
+                              className="text-foreground-muted"
+                            />
+                          )}
+                          <div>
+                            {part.state === 'input-streaming' ? 'Running ' : 'Ran '}
+                            <span className="text-foreground-lighter">{`${part.type.replace('tool-', '')}`}</span>
+                          </div>
+                        </div>
+                      )
+                    }
+                    case 'reasoning':
+                      return (
+                        <Reasoning
+                          key={`${message.id}-${index}}`}
+                          showReasoning={!!part.text}
+                          className={cn(
+                            'w-full dynamic-tool-item',
+                            '[&:not(.dynamic-tool-item+.dynamic-tool-item)]:mt-4 [&.dynamic-tool-item+.dynamic-tool-item]:mt-1 first:!mt-0',
+                            '[&:not(:has(+.dynamic-tool-item))]:mb-4'
+                          )}
+                          isStreaming={part.state === 'streaming'}
+                        >
+                          {part.text || part.text}
+                        </Reasoning>
+                      )
+                    case 'text':
+                      return (
+                        <ReactMarkdown
+                          key={`${id}-part-${index}`}
+                          className={cn(
+                            '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'
+                          )}
+                          remarkPlugins={[remarkGfm]}
+                          components={allMarkdownComponents}
+                        >
+                          {part.text}
+                        </ReactMarkdown>
+                      )

-        <div className="flex gap-4 w-auto overflow-hidden group">
-          {isUser && (
-            <ProfileImage
-              alt={profile?.username}
-              src={profile?.profileImageUrl}
-              className="w-5 h-5 shrink-0 rounded-full translate-y-0.5"
-            />
-          )}
+                    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

-          <div className="flex-1 min-w-0">
-            {shouldUsePartsRendering ? (
-              (() => {
-                return parts.map(
-                  (part: NonNullable<VercelMessage['parts']>[number], index: number) => {
-                    const isLastPart = index === parts.length - 1
-                    switch (part.type) {
-                      case 'dynamic-tool': {
+                      if (state === 'input-streaming') {
                        return (
                          <div
-                            key={`${id}-tool-${part.toolCallId}`}
-                            className={cn(
-                              'tool-item text-foreground-lighter flex items-center gap-2 py-2',
-                              '[&:not(.tool-item+.tool-item)]:mt-4 [&:not(:has(+.tool-item))]:mb-4',
-                              '[&:has(+.tool-item)]:border-b [&:has(+.tool-item)]:border-b-muted',
-                              'first:!mt-0 last:mb-0'
-                            )}
+                            key={`${id}-tool-loading-execute_sql`}
+                            className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2"
                          >
-                            {part.state === 'input-streaming' ? (
-                              <Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
-                            ) : (
-                              <CheckIcon
-                                strokeWidth={1.5}
-                                size={12}
-                                className="text-foreground-muted"
-                              />
-                            )}
-                            <div>
-                              {part.state === 'input-streaming' ? 'Running ' : 'Ran '}
-                              <span className="text-foreground-lighter">{`${part.toolName}`}</span>
-                            </div>
+                            <Loader2 className="w-4 h-4 animate-spin" />
+                            Writing SQL...
                          </div>
                        )
                      }
-                      case 'reasoning':
+
+                      if (state === 'input-available' || state === 'output-available') {
                        return (
-                          <Reasoning
-                            key={`${message.id}-${index}}`}
-                            showReasoning={!!part.text}
-                            className={cn(
-                              'w-full dynamic-tool-item',
-                              '[&:not(.dynamic-tool-item+.dynamic-tool-item)]:mt-4 [&.dynamic-tool-item+.dynamic-tool-item]:mt-1 first:!mt-0',
-                              '[&:not(:has(+.dynamic-tool-item))]:mb-4'
-                            )}
-                            isStreaming={part.state === 'streaming'}
+                          <div
+                            key={`${id}-tool-${toolCallId}`}
+                            className="w-auto overflow-x-hidden my-4 space-y-2"
                          >
-                            {part.text || part.text}
-                          </Reasoning>
-                        )
-                      case 'text':
-                        return (
-                          <ReactMarkdown
-                            key={`${id}-part-${index}`}
-                            className={cn(
-                              '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'
-                            )}
-                            remarkPlugins={[remarkGfm]}
-                            components={allMarkdownComponents}
-                          >
-                            {part.text}
-                          </ReactMarkdown>
-                        )
+                            <DisplayBlockRenderer
+                              messageId={id}
+                              toolCallId={toolCallId}
+                              initialArgs={{
+                                sql: (inputArgs?.sql as string) ?? '',
+                                label: inputArgs?.label,
+                                isWriteQuery: inputArgs?.isWriteQuery,
+                                view,
+                                xAxis,
+                                yAxis,
+                              }}
+                              initialResults={output}
+                              toolState={state}
+                              isLastPart={isLastPart}
+                              isLastMessage={isLastMessage}
+                              onResults={(args: {
+                                messageId: string
+                                resultId?: string
+                                results: unknown
+                              }) => {
+                                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 (
-                            <div
-                              key={`${id}-tool-loading-execute_sql`}
-                              className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2"
-                            >
-                              <Loader2 className="w-4 h-4 animate-spin" />
-                              Writing SQL...
-                            </div>
-                          )
-                        }
-
-                        if (state === 'input-available' || state === 'output-available') {
-                          return (
-                            <div
-                              key={`${id}-tool-${toolCallId}`}
-                              className="w-auto overflow-x-hidden my-4 space-y-2"
-                            >
-                              <DisplayBlockRenderer
-                                messageId={id}
-                                toolCallId={toolCallId}
-                                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 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)
-                                }}
-                              />
-                            </div>
-                          )
-                        }
-                        if (state === 'output-error') {
-                          return (
-                            <div key={`${id}-tool-${toolCallId}`} className="text-xs text-danger">
-                              Failed to execute SQL.
-                            </div>
-                          )
-                        }
-                        return null
-                      }
-                      case 'tool-deploy_edge_function': {
-                        const { toolCallId, state, input } = part
-                        if (state === 'input-streaming') {
-                          return (
-                            <div
-                              key={`${id}-tool-loading-deploy_edge_function`}
-                              className="my-4 rounded-lg border bg-surface-75 font-mono text-xs text-foreground-lighter py-2 px-3 flex items-center gap-2"
-                            >
-                              <Loader2 className="w-4 h-4 animate-spin" />
-                              {`Preparing Edge Function deployment...`}
-                            </div>
-                          )
-                        }
-                        if (state === 'input-available' || state === 'output-available') {
-                          return (
-                            <EdgeFunctionRenderer
-                              key={`${id}-tool-${toolCallId}`}
-                              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)
-                                }
+                                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 (
-                            <div key={`${id}-tool-${toolCallId}`} className="text-xs text-danger">
-                              Failed to deploy Edge Function.
-                            </div>
-                          )
-                        }
-                        return null
+                          </div>
+                        )
                      }
-
-                      case 'source-url':
-                      case 'source-document':
-                      case 'file':
-                        return null
-                      default:
-                        return null
+                      if (state === 'output-error') {
+                        return (
+                          <div key={`${id}-tool-${toolCallId}`} className="text-xs text-danger">
+                            Failed to execute SQL.
+                          </div>
+                        )
+                      }
+                      return null
+                    }
+                    case 'tool-deploy_edge_function': {
+                      const { toolCallId, state, input } = part
+                      if (state === 'input-streaming') {
+                        return (
+                          <div
+                            key={`${id}-tool-loading-execute_sql`}
+                            className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2"
+                          >
+                            <Loader2 className="w-4 h-4 animate-spin" />
+                            Writing Edge Function...
+                          </div>
+                        )
+                      }
+                      if (state === 'input-available' || state === 'output-available') {
+                        return (
+                          <EdgeFunctionRenderer
+                            key={`${id}-tool-${toolCallId}`}
+                            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') {
+                        return (
+                          <div key={`${id}-tool-${toolCallId}`} className="text-xs text-danger">
+                            Failed to deploy Edge Function.
+                          </div>
+                        )
+                      }
+                      return null
                    }
-                  }
-                )
-              })()
-            ) : hasTextContent ? (
-              <ReactMarkdown
-                className="prose prose-sm max-w-none break-words"
-                remarkPlugins={[remarkGfm]}
-                components={allMarkdownComponents}
-              >
-                {content}
-              </ReactMarkdown>
-            ) : (
-              <span className="text-foreground-lighter italic">Assistant is thinking...</span>
-            )}
-          </div>
-        </div>
-        {message.role === 'user' && (
-          <div className="flex items-center gap-4 mt-2 mb-1">
-            <span className="h-0.5 w-5 bg-muted" />
-            <div className="opacity-0 group-hover:opacity-100 transition-opacity">
-              <>
-                <ButtonTooltip
-                  type="text"
-                  icon={<Pencil size={14} strokeWidth={1.5} />}
-                  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',
-                    },
-                  }}
-                />

-                <ButtonTooltip
-                  type="text"
-                  icon={<Trash2 size={14} strokeWidth={1.5} />}
-                  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"
-                />
-              </>
-            </div>
-          </div>
-        )}
+                    case 'source-url':
+                    case 'source-document':
+                    case 'file':
+                      return null
+                    default:
+                      return null
+                  }
+                }
+              )
+            })()
+          ) : hasTextContent ? (
+            <ReactMarkdown
+              className="prose prose-sm max-w-none break-words"
+              remarkPlugins={[remarkGfm]}
+              components={allMarkdownComponents}
+            >
+              {content}
+            </ReactMarkdown>
+          ) : (
+            <span className="text-foreground-lighter italic">Assistant is thinking...</span>
+          )}
+        </div>
      </div>

+      {message.role === 'user' && (
+        <div className="flex items-center gap-4 mt-2 mb-1">
+          <span className="h-0.5 w-5 bg-muted" />
+          <div className="opacity-0 group-hover:opacity-100 transition-opacity">
+            <>
+              <ButtonTooltip
+                type="text"
+                icon={<Pencil size={14} strokeWidth={1.5} />}
+                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',
+                  },
+                }}
+              />
+
+              <ButtonTooltip
+                type="text"
+                icon={<Trash2 size={14} strokeWidth={1.5} />}
+                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"
+              />
+            </>
+          </div>
+        </div>
+      )}
      <DeleteMessageConfirmModal
        visible={showDeleteConfirmModal}
        onConfirm={() => {
@@ -391,6 +401,6 @@ export const Message = function Message({
        }}
        onCancel={() => setShowDeleteConfirmModal(false)}
      />
-    </MessageContext.Provider>
+    </div>
  )
 }
@@ -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 (
    <div className="w-auto overflow-x-hidden not-prose my-4 ">
      {language === 'edge' ? (
@@ -235,42 +184,29 @@ export const MarkdownPre = ({
      ) : language === 'sql' ? (
        readOnly ? (
          <CollapsibleCodeBlock value={cleanContent} language="sql" hideLineNumbers />
+        ) : isLoading ? (
+          <div className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2">
+            <Loader2 className="w-4 h-4 animate-spin" />
+            Writing SQL...
+          </div>
        ) : (
-          <DebouncedComponent
-            delay={isLoading ? 500 : 0}
-            value={cleanContent}
-            fallback={
-              <div className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2">
-                <Loader2 className="w-4 h-4 animate-spin" />
-                Writing SQL...
-              </div>
-            }
-          >
-            <QueryBlock
-              label={title}
-              sql={cleanContent}
-              chartConfig={{
-                type: 'bar',
-                cumulative: false,
-                xKey: xAxis ?? '',
-                yKey: yAxis ?? '',
-                view: isChart ? 'chart' : 'table',
-              }}
-              initialResults={results}
-              onRunQuery={onRunQuery}
-              onResults={onResultsReturned}
-              onUpdateChartConfig={({ chartConfig: config }) => {
-                chartConfig.current = { ...chartConfig.current, ...config }
-              }}
-              draggable={isDraggableToReports}
-              onDragStart={(e: DragEvent<Element>) => {
-                e.dataTransfer.setData(
-                  'application/json',
-                  JSON.stringify({ label: title, sql: cleanContent, config: chartConfig.current })
-                )
-              }}
-            />
-          </DebouncedComponent>
+          <DisplayBlockRenderer
+            messageId={id}
+            toolCallId={toolCallId}
+            initialArgs={{
+              sql: cleanContent,
+              label: title,
+              isWriteQuery: false,
+              view: isChart ? 'chart' : 'table',
+              xAxis: xAxis ?? '',
+              yAxis: yAxis ?? '',
+            }}
+            onError={() => {}}
+            showConfirmFooter={false}
+            onChartConfigChange={(config) => {
+              chartConfig.current = { ...config }
+            }}
+          />
        )
      ) : (
        <CodeBlock
@@ -1,17 +1,13 @@
 import dayjs from 'dayjs'
 import { Code, Play } from 'lucide-react'
-import { DragEvent, ReactNode, useCallback, useEffect, useMemo, useState } from 'react'
+import { DragEvent, ReactNode, useEffect, useMemo, useRef, useState } from 'react'
 import { Bar, BarChart, CartesianGrid, Cell, Tooltip, XAxis, YAxis } from 'recharts'
-import { toast } from 'sonner'

-import { useParams } from 'common'
 import { ReportBlockContainer } from 'components/interfaces/Reports/ReportBlock/ReportBlockContainer'
 import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartConfig'
 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 { Badge, ChartContainer, ChartTooltipContent, cn, CodeBlock, SQL_ICON } from 'ui'
+import { Badge, ChartContainer, ChartTooltipContent, cn, CodeBlock } from 'ui'
 import ShimmeringLoader from 'ui-patterns/ShimmeringLoader'
 import { ButtonTooltip } from '../ButtonTooltip'
 import { CHART_COLORS } from '../Charts/Charts.constants'
@@ -30,22 +26,22 @@ export const DEFAULT_CHART_CONFIG: ChartConfig = {
  view: 'table',
 }

-interface QueryBlockProps {
+export interface QueryBlockProps {
  id?: string
  label: string
  sql?: string
  isWriteQuery?: boolean
  chartConfig?: ChartConfig
  actions?: ReactNode
-  initialResults?: any[]
-  onResults?: (results: any[]) => 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<ChartConfig> }) => void
  draggable?: boolean
  onDragStart?: (e: DragEvent<Element>) => 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>(chartConfig)
  const { xKey, yKey, view = 'table' } = chartSettings

-  // Maintain local results state, seeded from initial results prop
-  const [localResults, setLocalResults] = useState<any[] | undefined>(initialResults)
-
-  const [showSql, setShowSql] = useState(!localResults)
-  const [readOnlyError, setReadOnlyError] = useState(false)
-  const [queryError, setQueryError] = useState<QueryResponseError>()
+  const [showSql, setShowSql] = useState(!results && !initialHideSql)
  const [focusDataIndex, setFocusDataIndex] = useState<number>()
+  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 (
    <ReportBlockContainer
@@ -187,16 +147,16 @@ export const QueryBlock = ({
                content: { side: 'bottom', text: showSql ? 'Hide query' : 'Show query' },
              }}
            />
-            {localResults && (
+            {hasResults && (
              <BlockViewConfiguration
                view={view}
                isChart={view === 'chart'}
                lockColumns={false}
                chartConfig={chartSettings}
-                columns={Object.keys(localResults?.[0] || {})}
-                changeView={(view) => {
-                  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={<Play size={14} strokeWidth={1.5} />}
              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 ? (
-                    <p>{`Query is running. You may cancel ongoing queries via the [SQL Editor](/project/${ref}/sql?viewOngoingQueries=true).`}</p>
-                  ) : (
-                    '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 && (
        <div
-          className={cn('shrink-0 w-full max-h-96 overflow-y-auto', {
-            'border-b': localResults !== undefined,
+          className={cn('h-full w-full max-h-full min-h-32 overflow-y-auto', {
+            'border-b': results !== undefined,
          })}
        >
          <CodeBlock
@@ -285,15 +228,15 @@ export const QueryBlock = ({
        </div>
      )}

-      {isExecuting && localResults === undefined && (
+      {isExecuting && !results && (
        <div className="p-3 w-full border-t">
          <ShimmeringLoader />
        </div>
      )}

-      {view === 'chart' && localResults !== undefined ? (
+      {view === 'chart' && results !== undefined ? (
        <>
-          {(localResults ?? []).length === 0 ? (
+          {(results ?? []).length === 0 ? (
            <div className="flex w-full h-full items-center justify-center py-3">
              <p className="text-foreground-light text-xs">No results returned from query</p>
            </div>
@@ -305,7 +248,7 @@ export const QueryBlock = ({
            <div className="flex-1 w-full">
              <ChartContainer
                className="aspect-auto px-3 py-2"
-                style={{ height: '250px', minHeight: '250px' }}
+                style={{ height: '230px', minHeight: '230px' }}
              >
                <BarChart
                  accessibilityLayer
@@ -350,14 +293,14 @@ export const QueryBlock = ({
        </>
      ) : (
        <>
-          {!isExecuting && !!queryError ? (
+          {!isExecuting && !!errorText ? (
            <div className={cn('flex-1 w-full overflow-auto relative border-t px-3.5 py-2')}>
-              <span className="font-mono text-xs">ERROR: {queryError.message}</span>
+              <span className="font-mono text-xs">ERROR: {errorText}</span>
            </div>
          ) : (
-            localResults && (
+            results && (
              <div className={cn('flex-1 w-full overflow-auto relative max-h-64')}>
-                <Results rows={localResults} />
+                <Results rows={results} />
              </div>
            )
          )}
@@ -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.
@@ -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
        })
@@ -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()
      }