all the ui changes

2026-10-13 21:29:08 +03:00 · 2025-09-15 14:57:57 +10:00
1 parent 8ae64d0a2c
commit 55f42c4eae
@@ -68,78 +68,22 @@ export const ReportBlock = ({
    <>
      {isSnippet ? (
        <QueryBlock
-          runQuery
-          isChart
-          draggable
          id={item.id}
-          isLoading={isLoading}
-          isRefreshing={isRefreshing}
          label={item.label}
          chartConfig={chartConfig}
          sql={sql}
-          maxHeight={232}
-          queryHeight={232}
          actions={
-            <ButtonTooltip
-              type="text"
-              icon={<X />}
-              className="w-7 h-7"
-              onClick={() => onRemoveChart({ metric: { key: item.attribute } })}
-              tooltip={{ content: { side: 'bottom', text: 'Remove chart' } }}
-            />
+            !isLoading && (
+              <ButtonTooltip
+                type="text"
+                icon={<X />}
+                className="w-7 h-7"
+                onClick={() => onRemoveChart({ metric: { key: item.attribute } })}
+                tooltip={{ content: { side: 'bottom', text: 'Remove chart' } }}
+              />
+            )
          }
          onUpdateChartConfig={onUpdateChart}
-          noResultPlaceholder={
-            <div
-              className={cn(
-                'flex flex-col gap-y-1 h-full w-full',
-                isLoading ? 'justify-start items-start p-2 gap-y-2' : 'justify-center px-4 gap-y-1'
-              )}
-            >
-              {isLoading ? (
-                <>
-                  <ShimmeringLoader className="w-full" />
-                  <ShimmeringLoader className="w-full w-3/4" />
-                  <ShimmeringLoader className="w-full w-1/2" />
-                </>
-              ) : isError ? (
-                <>
-                  <p className="text-xs text-foreground-light text-center">
-                    {snippetMissing ? 'SQL snippet cannot be found' : 'Error fetching SQL snippet'}
-                  </p>
-                  <p className="text-xs text-foreground-lighter text-center">
-                    {snippetMissing ? 'Please remove this block from your report' : error.message}
-                  </p>
-                </>
-              ) : (
-                <>
-                  <p className="text-xs text-foreground-light text-center">
-                    No results returned from query
-                  </p>
-                  <p className="text-xs text-foreground-lighter text-center">
-                    Results from the SQL query can be viewed as a table or chart here
-                  </p>
-                </>
-              )}
-            </div>
-          }
-          readOnlyErrorPlaceholder={
-            <div className="flex flex-col h-full justify-center items-center text-center">
-              <p className="text-xs text-foreground-light">
-                SQL query is not read-only and cannot be rendered
-              </p>
-              <p className="text-xs text-foreground-lighter text-center">
-                Queries that involve any mutation will not be run in reports
-              </p>
-              <Button
-                type="default"
-                className="mt-2"
-                onClick={() => onRemoveChart({ metric: { key: item.attribute } })}
-              >
-                Remove chart
-              </Button>
-            </div>
-          }
        />
      ) : isDeprecatedChart ? (
        <DeprecatedChartBlock
@@ -1,10 +1,10 @@
-import { GripHorizontal, Loader2 } from 'lucide-react'
+import { Code, GripHorizontal, Loader2 } from 'lucide-react'
 import { DragEvent, PropsWithChildren, ReactNode } from 'react'
 import { cn, Tooltip, TooltipContent, TooltipTrigger } from 'ui'

 interface ReportBlockContainerProps {
-  icon: ReactNode
-  label: string
+  icon?: ReactNode
+  label: string | ReactNode
  actions: ReactNode
  loading?: boolean
  draggable?: boolean
@@ -39,33 +39,18 @@ export const ReportBlockContainer = ({
        <TooltipTrigger asChild>
          <div
            className={cn(
-              'grid-item-drag-handle flex py-1 pl-3 pr-1 items-center gap-2 z-10 shrink-0 group',
+              'grid-item-drag-handle flex py-1 pl-3 pr-1 items-center gap-2 z-10 shrink-0 group h-9',
              draggable && 'cursor-move'
            )}
          >
-            <div
-              className={cn(
-                showDragHandle && 'transition-opacity opacity-100 group-hover:opacity-0'
-              )}
-            >
-              {loading ? (
-                <Loader2
-                  size={(icon as any)?.props?.size ?? 16}
-                  className="text-foreground-lighter animate-spin"
-                />
-              ) : (
-                icon
-              )}
-            </div>
-            {showDragHandle && (
+            {showDragHandle ? (
              <div className="absolute left-3 top-2.5 z-10 opacity-0 transition-opacity group-hover:opacity-100">
                <GripHorizontal size={16} strokeWidth={1.5} />
              </div>
+            ) : (
+              <Code size={16} strokeWidth={1.5} className="text-foreground-muted" />
            )}
-            <h3
-              title={label}
-              className="!text-xs font-medium text-foreground-light flex-1 truncate"
-            >
+            <h3 title={label} className="heading-meta flex-1">
              {label}
            </h3>
            <div className="flex items-center">{actions}</div>
@@ -37,7 +37,7 @@ import {
  ConversationContent,
  ConversationScrollButton,
 } from './elements/Conversation'
-import { MemoizedMessage } from './Message'
+import { Message } from './Message'

 interface AIAssistantProps {
  initialMessages?: MessageType[] | undefined
@@ -113,7 +113,12 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
        snap.saveMessage([lastUserMessageRef.current, message])
        lastUserMessageRef.current = null
      } else {
-        snap.saveMessage(message)
+        const existingMessage = snap.activeChat?.messages.find((m) => m.id === message.id)
+        if (existingMessage) {
+          snap.updateMessage(message)
+        } else {
+          snap.saveMessage(message)
+        }
      }
    },
    [snap]
@@ -122,6 +127,17 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
  // TODO(refactor): This useChat hook should be moved down into each chat session.
  // That way we won't have to disable switching chats while the chat is loading,
  // and don't run the risk of messages getting mixed up between chats.
+  // Sanitize messages to remove Valtio proxy wrappers that can't be cloned
+  const sanitizedMessages = useMemo(() => {
+    if (!snap.activeChat?.messages) return undefined
+
+    return snap.activeChat.messages.map((msg: any) => {
+      // Convert proxy objects to plain objects
+      const plainMessage = JSON.parse(JSON.stringify(msg))
+      return plainMessage
+    })
+  }, [snap.activeChat?.messages])
+
  const {
    messages: chatMessages,
    status: chatStatus,
@@ -133,26 +149,31 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
    regenerate,
  } = useChat({
    id: snap.activeChatId,
-    // [Alaister] typecast is needed here because valtio returns readonly arrays
-    // and useChat expects a mutable array
    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
-    messages: snap.activeChat?.messages as unknown as MessageType[] | undefined,
+    messages: sanitizedMessages,
    async onToolCall({ toolCall }) {
+      if (toolCall.dynamic) {
+        return
+      }
+
      if (toolCall.toolName === 'rename_chat') {
        const { newName } = toolCall.input as { newName: string }
+
        if (snap.activeChatId && newName?.trim()) {
          snap.renameChat(snap.activeChatId, newName.trim())
+
          addToolResult({
            tool: toolCall.toolName,
            toolCallId: toolCall.toolCallId,
            output: 'Chat renamed',
          })
+        } else {
+          addToolResult({
+            tool: toolCall.toolName,
+            toolCallId: toolCall.toolCallId,
+            output: 'Failed to rename chat: Invalid chat or name',
+          })
        }
-        addToolResult({
-          tool: toolCall.toolName,
-          toolCallId: toolCall.toolCallId,
-          output: 'Failed to rename chat: Invalid chat or name',
-        })
      }
    },
    transport: new DefaultChatTransport({
@@ -200,11 +221,10 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
  const isChatLoading = chatStatus === 'submitted' || chatStatus === 'streaming'

  const updateMessage = useCallback(
-    (_args: { messageId: string; resultId?: string; results: any[] }) => {
-      // Intentionally no-op: we are not saving results into chat state
-      return
+    (updatedMessage: MessageType) => {
+      snap.updateMessage(updatedMessage)
    },
-    []
+    [snap]
  )

  const deleteMessageFromHere = useCallback(
@@ -266,25 +286,23 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
        const isAfterEditedMessage = editingMessageId
          ? chatMessages.findIndex((m) => m.id === editingMessageId) < index
          : false
+        const isLastMessage = index === chatMessages.length - 1

        return (
-          <MemoizedMessage
+          <Message
+            id={message.id}
            key={message.id}
            message={message}
-            status={chatStatus}
+            isLoading={chatStatus === 'submitted' || chatStatus === 'streaming'}
+            readOnly={message.role === 'user'}
            onResults={updateMessage}
-            addToolResult={
-              addToolResult as unknown as (args: {
-                tool: string
-                toolCallId: string
-                output: unknown
-              }) => Promise<void>
-            }
+            addToolResult={addToolResult}
            onDelete={deleteMessageFromHere}
            onEdit={editMessage}
            isAfterEditedMessage={isAfterEditedMessage}
            isBeingEdited={isBeingEdited}
            onCancelEdit={cancelEdit}
+            isLastMessage={isLastMessage}
          />
        )
      }),
@@ -444,7 +462,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
                <motion.span
                  animate={{ opacity: [1, 0] }}
                  transition={{ duration: 1, repeat: Infinity, ease: 'linear' }}
-                  className="inline-block w-1.5 h-4 bg-foreground-lighter"
+                  className="inline-block w-1.5 h-4 bg-foreground-lighter mt-4"
                />
              )}
            </ConversationContent>
@@ -19,13 +19,13 @@ export const ConfirmFooter = ({
  onConfirm,
 }: PropsWithChildren<ConfirmFooterProps>) => {
  return (
-    <div className="rounded-lg border bg-surface-75 text-xs text-foreground-lighter p-2 pl-3 flex items-center justify-between gap-3">
+    <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">
        <Button size="tiny" type="outline" onClick={onCancel} disabled={isLoading}>
          {cancelLabel}
        </Button>
-        <Button size="tiny" type="outline" onClick={onConfirm} disabled={isLoading}>
+        <Button size="tiny" type="primary" onClick={onConfirm} disabled={isLoading}>
          {isLoading ? 'Working…' : confirmLabel}
        </Button>
      </div>
@@ -1,7 +1,6 @@
 import { PermissionAction } from '@supabase/shared-types/out/constants'
-import type { UIDataTypes, UIMessagePart, UITools } from 'ai'
 import { useRouter } from 'next/router'
-import { DragEvent, PropsWithChildren, useMemo, useState } from 'react'
+import { DragEvent, PropsWithChildren, useEffect, useState } from 'react'

 import { useParams } from 'common'
 import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartConfig'
@@ -9,55 +8,47 @@ import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
 import { useAsyncCheckProjectPermissions } from 'hooks/misc/useCheckPermissions'
 import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization'
 import { useProfile } from 'lib/profile'
-import { useAiAssistantStateSnapshot } from 'state/ai-assistant-state'
-import { Badge } from 'ui'
+
+import { ConfirmFooter } from './ConfirmFooter'
 import { DEFAULT_CHART_CONFIG, QueryBlock } from '../QueryBlock/QueryBlock'
 import { identifyQueryType } from './AIAssistant.utils'
-import { findResultForManualId } from './Message.utils'
+import { usePrimaryDatabase } from 'data/read-replicas/replicas-query'
+import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'

 interface DisplayBlockRendererProps {
  messageId: string
  toolCallId: string
-  manualId?: string
  initialArgs: {
    sql: string
    label?: string
+    isWriteQuery?: boolean
    view?: 'table' | 'chart'
    xAxis?: string
    yAxis?: string
-    runQuery?: boolean
  }
-  messageParts: UIMessagePart<UIDataTypes, UITools>[] | undefined
-  isLoading: boolean
-  onResults: (args: { messageId: string; resultId?: string; results: any[] }) => void
+  results?: unknown
+  onResults: (args: { messageId: string; resultId?: string; results: unknown }) => void
  onError?: (args: { messageId: string; resultId?: string; errorText: string }) => void
-  /**
-   * External signal to trigger a query run.
-   * Increment this number to trigger another run. The component will
-   * translate this into a boolean toggle for QueryBlock.runQuery.
-   */
-  triggerRunSignal?: number
-  /** Notify parent when a run starts or ends (for shared loading states) */
-  onRunStateChange?: (isRunning: boolean) => void
+  toolState?: 'input-streaming' | 'input-available' | 'output-available' | 'output-error'
+  isLastPart?: boolean
+  isLastMessage?: boolean
 }

 export const DisplayBlockRenderer = ({
  messageId,
  toolCallId,
-  manualId,
  initialArgs,
-  messageParts,
-  isLoading,
+  results,
  onResults,
  onError,
-  triggerRunSignal = 0,
-  onRunStateChange,
+  toolState,
+  isLastPart = false,
+  isLastMessage = false,
 }: PropsWithChildren<DisplayBlockRendererProps>) => {
  const router = useRouter()
  const { ref } = useParams()
  const { profile } = useProfile()
  const { data: org } = useSelectedOrganizationQuery()
-  const snap = useAiAssistantStateSnapshot()

  const { mutate: sendEvent } = useSendEventMutation()
  const { can: canCreateSQLSnippet } = useAsyncCheckProjectPermissions(
@@ -76,21 +67,28 @@ export const DisplayBlockRenderer = ({
    yKey: initialArgs.yAxis ?? '',
  }))

-  const isChart = initialArgs.view === 'chart'
-  const resultId = manualId || toolCallId
-  const liveResultData = useMemo(
-    () => (manualId ? findResultForManualId(messageParts, manualId) : undefined),
-    [messageParts, manualId]
-  )
-  const cachedResults = useMemo(
-    () => snap.getCachedSQLResults({ messageId, snippetId: resultId }),
-    [snap, messageId, resultId]
-  )
-  const displayData = liveResultData ?? cachedResults
+  const resultId = toolCallId
+  const displayData = Array.isArray(results) ? results : undefined
  const isDraggableToReports = canCreateSQLSnippet && router.pathname.endsWith('/reports/[id]')
  const label = initialArgs.label || 'SQL Results'
+  const isWriteQuery = initialArgs.isWriteQuery || false
  const sqlQuery = initialArgs.sql
-  const runQueryFlag = (triggerRunSignal || 0) % 2 === 1
+  const [isRunning, setIsRunning] = useState(false)
+
+  const { database: primaryDatabase } = usePrimaryDatabase({ projectRef: ref })
+  const readOnlyConnectionString = primaryDatabase?.connection_string_read_only
+  const postgresConnectionString = primaryDatabase?.connectionString
+
+  const { mutate: execute } = useExecuteSqlMutation({
+    onSuccess: (data) => {
+      onResults({ messageId, resultId, results: data.result })
+      setIsRunning(false)
+    },
+    onError: (error) => {
+      onError?.({ messageId, resultId, errorText: error.message })
+      setIsRunning(false)
+    },
+  })

  const handleRunQuery = (queryType: 'select' | 'mutation') => {
    sendEvent({
@@ -104,7 +102,14 @@ export const DisplayBlockRenderer = ({
        organization: org?.slug ?? 'Unknown',
      },
    })
-    onRunStateChange?.(true)
+  }
+
+  const handleExecute = (queryType: 'select' | 'mutation') => {
+    handleRunQuery(queryType)
+    if (ref && postgresConnectionString) {
+      setIsRunning(true)
+      execute({ projectRef: ref, connectionString: postgresConnectionString, sql: sqlQuery })
+    }
  }

  const handleUpdateChartConfig = ({
@@ -122,42 +127,51 @@ export const DisplayBlockRenderer = ({
    )
  }

+  const resolvedHasDecision = !!results
+  const shouldShowConfirmFooter =
+    !resolvedHasDecision && toolState === 'input-available' && isLastPart && isLastMessage
+
  return (
-    <div className="display-block w-auto overflow-x-hidden !my-6">
+    <div className="display-block w-auto overflow-x-hidden">
      <QueryBlock
        label={label}
+        isWriteQuery={isWriteQuery}
        sql={sqlQuery}
-        lockColumns={true}
-        showSql={!isChart}
-        results={displayData}
+        initialResults={displayData}
        chartConfig={chartConfig}
-        isChart={isChart}
-        showRunButtonIfNotReadOnly={true}
-        isLoading={isLoading}
-        draggable={isDraggableToReports}
-        runQuery={runQueryFlag}
-        tooltip={
-          isDraggableToReports ? (
-            <div className="flex items-center gap-x-2">
-              <Badge variant="success" className="text-xs rounded px-1">
-                NEW
-              </Badge>
-              <p>Drag to add this chart into your custom report</p>
-            </div>
-          ) : undefined
-        }
        onResults={(results) => {
          onResults({ messageId, resultId, results })
-          onRunStateChange?.(false)
+          setIsRunning(false)
        }}
        onRunQuery={handleRunQuery}
        onUpdateChartConfig={handleUpdateChartConfig}
        onError={(errorText) => {
          onError?.({ messageId, resultId, errorText })
-          onRunStateChange?.(false)
+          setIsRunning(false)
        }}
+        draggable={isDraggableToReports}
        onDragStart={handleDragStart}
+        disabled={!resolvedHasDecision || isRunning}
+        isExternallyExecuting={isRunning}
      />
+
+      {shouldShowConfirmFooter && (
+        <div className="mx-4">
+          <ConfirmFooter
+            message="Assistant wants to run this query"
+            cancelLabel="Skip"
+            confirmLabel={isRunning ? 'Running…' : 'Run Query'}
+            isLoading={isRunning}
+            onCancel={async () => {
+              setIsRunning(false)
+              onResults({ messageId, resultId, results: 'User skipped running the query' })
+            }}
+            onConfirm={() => {
+              handleExecute(isWriteQuery ? 'mutation' : 'select')
+            }}
+          />
+        </div>
+      )}
    </div>
  )
 }
@@ -0,0 +1,79 @@
+import { PropsWithChildren, useState } from 'react'
+import { toast } from 'sonner'
+
+import { useParams } from 'common'
+import { useEdgeFunctionDeployMutation } from 'data/edge-functions/edge-functions-deploy-mutation'
+import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
+import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization'
+import { EdgeFunctionBlock } from '../EdgeFunctionBlock/EdgeFunctionBlock'
+import { ConfirmFooter } from './ConfirmFooter'
+
+interface EdgeFunctionRendererProps {
+  label: string
+  code: string
+  functionName: string
+  onDeployed?: (result: { success: true } | { success: false; errorText: string }) => void
+}
+
+export const EdgeFunctionRenderer = ({
+  label,
+  code,
+  functionName,
+  onDeployed,
+}: PropsWithChildren<EdgeFunctionRendererProps>) => {
+  const { ref } = useParams()
+  const { data: org } = useSelectedOrganizationQuery()
+  const { mutate: sendEvent } = useSendEventMutation()
+  const [isRunning, setIsRunning] = useState(false)
+
+  const { mutateAsync: deployFunction } = useEdgeFunctionDeployMutation({
+    onSuccess: () => {
+      toast.success('Successfully deployed edge function')
+    },
+  })
+
+  const handleDeploy = async () => {
+    if (!ref || !functionName) return
+    try {
+      setIsRunning(true)
+      await deployFunction({
+        projectRef: ref,
+        slug: functionName,
+        metadata: {
+          entrypoint_path: 'index.ts',
+          name: functionName,
+          verify_jwt: true,
+        },
+        files: [{ name: 'index.ts', content: code }],
+      })
+      sendEvent({
+        action: 'edge_function_deploy_button_clicked',
+        properties: { origin: 'functions_ai_assistant' },
+        groups: { project: ref ?? 'Unknown', organization: org?.slug ?? 'Unknown' },
+      })
+      onDeployed?.({ success: true })
+    } catch (e: any) {
+      const errMsg = e?.message ?? 'Unknown error'
+      toast.error(`Failed to deploy function: ${errMsg}`)
+      onDeployed?.({ success: false, errorText: errMsg })
+    } finally {
+      setIsRunning(false)
+    }
+  }
+
+  return (
+    <div className="w-auto overflow-x-hidden my-4">
+      <EdgeFunctionBlock label={label} code={code} functionName={functionName} hideDeployButton />
+      <div className="mt-2">
+        <ConfirmFooter
+          message="Assistant wants to deploy this EdgeFunction"
+          cancelLabel="Skip"
+          confirmLabel={isRunning ? 'Deploying…' : 'Deploy'}
+          isLoading={isRunning}
+          onCancel={() => onDeployed?.({ success: false, errorText: 'Skipped' })}
+          onConfirm={handleDeploy}
+        />
+      </div>
+    </div>
+  )
+}
@@ -14,9 +14,11 @@ import { useProfile } from 'lib/profile'
 import { cn, markdownComponents, WarningIcon } from 'ui'
 import { ButtonTooltip } from '../ButtonTooltip'
 import { EdgeFunctionBlock } from '../EdgeFunctionBlock/EdgeFunctionBlock'
+import { EdgeFunctionRenderer } from './EdgeFunctionRenderer'
+import { ConfirmFooter } from './ConfirmFooter'
 import { DeleteMessageConfirmModal } from './DeleteMessageConfirmModal'
 import { DisplayBlockRenderer } from './DisplayBlockRenderer'
-import { ConfirmFooter } from './ConfirmFooter'
+
 import {
  Heading3,
  Hyperlink,
@@ -26,6 +28,7 @@ import {
  OrderedList,
 } from './MessageMarkdown'
 import { Reasoning } from './elements/Reasoning'
+import type { AssistantMessageType } from 'state/ai-assistant-state'

 interface MessageContextType {
  isLoading: boolean
@@ -47,27 +50,19 @@ interface MessageProps {
  message: VercelMessage
  isLoading: boolean
  readOnly?: boolean
-  status?: string
  action?: ReactNode
  variant?: 'default' | 'warning'
-  onResults: ({
-    messageId,
-    resultId,
-    results,
-  }: {
-    messageId: string
-    resultId?: string
-    results: any[]
-  }) => void
+  onResults: (message: VercelMessage) => void
  addToolResult?: (args: { tool: string; toolCallId: string; output: unknown }) => Promise<void>
  onDelete: (id: string) => void
  onEdit: (id: string) => void
  isAfterEditedMessage: boolean
  isBeingEdited: boolean
  onCancelEdit: () => void
+  isLastMessage?: boolean
 }

-const Message = function Message({
+export const Message = function Message({
  id,
  message,
  isLoading,
@@ -80,35 +75,35 @@ const Message = function Message({
  onEdit,
  isAfterEditedMessage = false,
  isBeingEdited = false,
-  status,
  onCancelEdit,
+  isLastMessage = false,
 }: PropsWithChildren<MessageProps>) {
  const { profile } = useProfile()
-  const { ref } = useParams()
-  const { data: project } = useSelectedProjectQuery()
-  const { mutateAsync: deployFunction } = useEdgeFunctionDeployMutation()
  const [showDeleteConfirmModal, setShowDeleteConfirmModal] = useState(false)
-  const [toolRunSignal, setToolRunSignal] = useState<Record<string, number>>({})
-  const [toolRunning, setToolRunning] = useState<Record<string, boolean>>({})
  const allMarkdownComponents: Partial<Components> = useMemo(
    () => ({
      ...markdownComponents,
      ...baseMarkdownComponents,
      pre: ({ children }) => (
-        <MarkdownPre id={id} onResults={onResults}>
+        <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)
+          }}
+        >
          {children}
        </MarkdownPre>
      ),
    }),
-    [id, onResults]
+    [id, onResults, message]
  )
-  const safeStringify = (value: unknown) => {
-    try {
-      return JSON.stringify(value)
-    } catch (_e) {
-      return 'null'
-    }
-  }

  if (!message) {
    console.error(`Message component received undefined message prop for id: ${id}`)
@@ -130,7 +125,7 @@ const Message = function Message({
    <MessageContext.Provider value={{ isLoading, readOnly }}>
      <div
        className={cn(
-          'text-foreground-light text-sm first:mt-0',
+          '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'
@@ -155,15 +150,17 @@ const Message = function Message({
              (() => {
                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(
-                              'border rounded-md border-muted heading-meta flex items-center gap-2 text-foreground-lighter py-2 px-3 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'
+                              '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' ? (
@@ -175,7 +172,10 @@ const Message = function Message({
                                className="text-foreground-muted"
                              />
                            )}
-                            {`${part.toolName}`}
+                            <div>
+                              {part.state === 'input-streaming' ? 'Running ' : 'Ran '}
+                              <span className="text-foreground-lighter">{`${part.toolName}`}</span>
+                            </div>
                          </div>
                        )
                      }
@@ -199,7 +199,7 @@ const Message = function Message({
                          <ReactMarkdown
                            key={`${id}-part-${index}`}
                            className={cn(
-                              'max-w-none prose prose-sm [&>div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h2:text-lg prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground prose-ol:space-y-3 prose-ul:space-y-3 prose-li:my-0 break-words [&>p:not(:last-child)]:!mb-2 [&>*>p:first-child]:!mt-0 [&>*>p:last-child]:!mb-0 [&>*>*>p:first-child]:!mt-0 [&>*>*>p:last-child]:!mb-0 [&>ol>li]:!pl-4',
+                              'max-w-none space-y-4 prose prose-sm prose-li:mt-1 [&>div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h2:text-lg prose-h3:no-underline prose-h3:text-base prose-h3:mb-4 prose-strong:font-medium prose-strong:text-foreground prose-ol:space-y-3 prose-ul:space-y-3 prose-li:my-0 break-words [&>p:not(:last-child)]:!mb-2 [&>*>p:first-child]:!mt-0 [&>*>p:last-child]:!mb-0 [&>*>*>p:first-child]:!mt-0 [&>*>*>p:last-child]:!mb-0 [&>ol>li]:!pl-4',
                              isUser && 'text-foreground [&>p]:font-medium',
                              isBeingEdited && 'animate-pulse'
                            )}
@@ -210,45 +210,17 @@ const Message = function Message({
                          </ReactMarkdown>
                        )

-                      case 'tool-display_query': {
-                        const { toolCallId, state, input } = part
-                        if (state === 'input-streaming' || state === 'input-available') {
-                          return (
-                            <div
-                              key={`${id}-tool-loading-display_query`}
-                              className="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" />
-                              {`Calling display_query...`}
-                            </div>
-                          )
-                        }
-                        if (state === 'output-available') {
-                          return (
-                            <DisplayBlockRenderer
-                              key={`${id}-tool-${toolCallId}`}
-                              messageId={id}
-                              toolCallId={toolCallId}
-                              manualId={(input as any).manualToolCallId}
-                              initialArgs={input as any}
-                              messageParts={parts}
-                              isLoading={false}
-                              onResults={onResults}
-                            />
-                          )
-                        }
-                        return null
-                      }
                      case 'tool-execute_sql': {
-                        const { toolCallId, state, input } = part
+                        const { toolCallId, state, input, output } = part
+
                        if (state === 'input-streaming') {
                          return (
                            <div
                              key={`${id}-tool-loading-execute_sql`}
-                              className="rounded-lg border bg-surface-75 font-mono text-xs text-foreground-lighter py-2 px-3 flex items-center gap-2"
+                              className="my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-2"
                            >
                              <Loader2 className="w-4 h-4 animate-spin" />
-                              {`Preparing SQL execution...`}
+                              {'Writing: ' + (input as any)?.label || 'SQL Query'}
                            </div>
                          )
                        }
@@ -262,14 +234,27 @@ const Message = function Message({
                              <DisplayBlockRenderer
                                messageId={id}
                                toolCallId={toolCallId}
-                                initialArgs={{ sql: (input as any).sql, label: 'SQL Query' }}
-                                messageParts={parts}
-                                isLoading={!!toolRunning[toolCallId]}
-                                onResults={({ results }) => {
-                                  addToolResult?.({
+                                initialArgs={{
+                                  sql: (input as any).sql,
+                                  label: (input as any).label,
+                                  isWriteQuery: (input as any).isWriteQuery,
+                                }}
+                                results={output}
+                                toolState={state}
+                                isLastPart={isLastPart}
+                                isLastMessage={isLastMessage}
+                                onResults={async (args: {
+                                  messageId: string
+                                  resultId?: string
+                                  results: unknown
+                                }) => {
+                                  const resultId = args.resultId
+                                  const results = args.results as any[]
+
+                                  await addToolResult?.({
                                    tool: 'execute_sql',
                                    toolCallId: String(toolCallId),
-                                    output: safeStringify(results),
+                                    output: results,
                                  })
                                }}
                                onError={async ({ errorText }) => {
@@ -278,38 +263,9 @@ const Message = function Message({
                                    toolCallId: String(toolCallId),
                                    output: `Error: ${errorText}`,
                                  })
+                                  onResults(message)
                                }}
-                                triggerRunSignal={toolRunSignal[toolCallId] || 0}
-                                onRunStateChange={(isRunning) =>
-                                  setToolRunning((prev) => ({ ...prev, [toolCallId]: isRunning }))
-                                }
                              />
-
-                              {state === 'input-available' && (
-                                <div className="!-mt-5">
-                                  <ConfirmFooter
-                                    message="Assistant wants to run this query"
-                                    cancelLabel="Skip"
-                                    confirmLabel="Run Query"
-                                    isLoading={!!toolRunning[toolCallId]}
-                                    onCancel={async () => {
-                                      await addToolResult?.({
-                                        tool: 'execute_sql',
-                                        toolCallId: String(toolCallId),
-                                        output: 'The user skipped running the query',
-                                      })
-                                      setToolRunning((prev) => ({ ...prev, [toolCallId]: false }))
-                                    }}
-                                    onConfirm={() => {
-                                      setToolRunning((prev) => ({ ...prev, [toolCallId]: true }))
-                                      setToolRunSignal((prev) => ({
-                                        ...prev,
-                                        [toolCallId]: (prev[toolCallId] || 0) + 1,
-                                      }))
-                                    }}
-                                  />
-                                </div>
-                              )}
                            </div>
                          )
                        }
@@ -337,71 +293,29 @@ const Message = function Message({
                        }
                        if (state === 'input-available' || state === 'output-available') {
                          return (
-                            <div
+                            <EdgeFunctionRenderer
                              key={`${id}-tool-${toolCallId}`}
-                              className="w-auto overflow-x-hidden my-4"
-                            >
-                              <EdgeFunctionBlock
-                                label={(input as any).name || 'Edge Function'}
-                                code={(input as any).code}
-                                functionName={(input as any).name || 'my-function'}
-                                showCode={!readOnly}
-                                hideDeployButton
-                              />
-                              {state === 'input-available' && (
-                                <div className="mt-2">
-                                  <ConfirmFooter
-                                    message="Assistant wants to deploy this EdgeFunction"
-                                    cancelLabel="Skip"
-                                    confirmLabel={toolRunning[toolCallId] ? 'Deploying…' : 'Deploy'}
-                                    isLoading={!!toolRunning[toolCallId]}
-                                    onCancel={async () => {
-                                      await addToolResult?.({
-                                        tool: 'deploy_edge_function',
-                                        toolCallId: String(toolCallId),
-                                        output: 'Skipped',
-                                      })
-                                    }}
-                                    onConfirm={async () => {
-                                      try {
-                                        setToolRunning((prev) => ({
-                                          ...prev,
-                                          [toolCallId]: true,
-                                        }))
-                                        await deployFunction({
-                                          projectRef: ref!,
-                                          slug: (input as any).name || 'my-function',
-                                          metadata: {
-                                            entrypoint_path: 'index.ts',
-                                            name: (input as any).name || 'my-function',
-                                            verify_jwt: true,
-                                          },
-                                          files: [
-                                            { name: 'index.ts', content: (input as any).code },
-                                          ],
-                                        })
-                                        await addToolResult?.({
-                                          tool: 'deploy_edge_function',
-                                          toolCallId: String(toolCallId),
-                                          output: 'Deployed successfully',
-                                        })
-                                      } catch (e: any) {
-                                        await addToolResult?.({
-                                          tool: 'deploy_edge_function',
-                                          toolCallId: String(toolCallId),
-                                          output: `Error: ${e?.message ?? String(e)}`,
-                                        })
-                                      } finally {
-                                        setToolRunning((prev) => ({
-                                          ...prev,
-                                          [toolCallId]: false,
-                                        }))
-                                      }
-                                    }}
-                                  />
-                                </div>
-                              )}
-                            </div>
+                              label={(input as any).name || 'Edge Function'}
+                              code={(input as any).code}
+                              functionName={(input as any).name || 'my-function'}
+                              onDeployed={async (res) => {
+                                if (res.success) {
+                                  await addToolResult?.({
+                                    tool: 'deploy_edge_function',
+                                    toolCallId: String(toolCallId),
+                                    output: 'Deployed successfully',
+                                  })
+                                  onResults(message)
+                                } else {
+                                  await addToolResult?.({
+                                    tool: 'deploy_edge_function',
+                                    toolCallId: String(toolCallId),
+                                    output: res.errorText,
+                                  })
+                                  onResults(message)
+                                }
+                              }}
+                            />
                          )
                        }
                        if (state === 'output-error') {
@@ -413,36 +327,7 @@ const Message = function Message({
                        }
                        return null
                      }
-                      case 'tool-display_edge_function': {
-                        const { toolCallId, state, input } = part
-                        if (state === 'input-streaming' || state === 'input-available') {
-                          return (
-                            <div
-                              key={`${id}-tool-loading-display_edge_function`}
-                              className="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" />
-                              {`Calling display_edge_function...`}
-                            </div>
-                          )
-                        }
-                        if (state === 'output-available') {
-                          return (
-                            <div
-                              key={`${id}-tool-${toolCallId}`}
-                              className="w-auto overflow-x-hidden my-4"
-                            >
-                              <EdgeFunctionBlock
-                                label={(input as any).name || 'Edge Function'}
-                                code={(input as any).code}
-                                functionName={(input as any).name || 'my-function'}
-                                showCode={!readOnly}
-                              />
-                            </div>
-                          )
-                        }
-                        return null
-                      }
+
                      case 'source-url':
                      case 'source-document':
                      case 'file':
@@ -464,44 +349,43 @@ const Message = function Message({
            ) : (
              <span className="text-foreground-lighter italic">Assistant is thinking...</span>
            )}
-
-            {/* Action button - only show for user messages on hover */}
-            <div className="opacity-0 group-hover:opacity-100 transition-opacity mt-1 mb-2">
-              {message.role === 'user' && (
-                <>
-                  <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>
        </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>
+        )}
      </div>

      <DeleteMessageConfirmModal
@@ -516,54 +400,3 @@ const Message = function Message({
    </MessageContext.Provider>
  )
 }
-
-export const MemoizedMessage = memo(
-  ({
-    message,
-    status,
-    onResults,
-    addToolResult,
-    onDelete,
-    onEdit,
-    isAfterEditedMessage,
-    isBeingEdited,
-    onCancelEdit,
-  }: {
-    message: VercelMessage
-    status: string
-    onResults: ({
-      messageId,
-      resultId,
-      results,
-    }: {
-      messageId: string
-      resultId?: string
-      results: any[]
-    }) => void
-    addToolResult?: (args: { tool: string; toolCallId: string; output: unknown }) => Promise<void>
-    onDelete: (id: string) => void
-    onEdit: (id: string) => void
-    isAfterEditedMessage: boolean
-    isBeingEdited: boolean
-    onCancelEdit: () => void
-  }) => {
-    return (
-      <Message
-        id={message.id}
-        message={message}
-        readOnly={message.role === 'user'}
-        isLoading={status === 'submitted' || status === 'streaming'}
-        status={status}
-        onResults={onResults}
-        addToolResult={addToolResult}
-        onDelete={onDelete}
-        onEdit={onEdit}
-        isAfterEditedMessage={isAfterEditedMessage}
-        isBeingEdited={isBeingEdited}
-        onCancelEdit={onCancelEdit}
-      />
-    )
-  }
-)
-
-MemoizedMessage.displayName = 'MemoizedMessage'
@@ -164,14 +164,12 @@ const MemoizedQueryBlock = memo(
      delay={isLoading ? 500 : 0}
      value={sql}
      fallback={
-        <div className="bg-surface-100 border-overlay rounded border shadow-sm px-3 py-2 text-xs">
+        <div className="heading-meta h-9 bg-surface-75 border-overlay rounded border shadow-sm px-3">
          Writing SQL...
        </div>
      }
    >
      <QueryBlock
-        lockColumns
-        showRunButtonIfNotReadOnly
        label={title}
        sql={sql}
        chartConfig={{
@@ -181,26 +179,12 @@ const MemoizedQueryBlock = memo(
          yKey: yAxis ?? '',
          view: isChart ? 'chart' : 'table',
        }}
-        tooltip={
-          isDraggable ? (
-            <div className="flex items-center gap-x-2">
-              <Badge variant="success" className="text-xs rounded px-1">
-                NEW
-              </Badge>
-              <p>Drag to add this chart into your custom report</p>
-            </div>
-          ) : undefined
-        }
-        showSql={!isChart}
-        isChart={isChart}
-        isLoading={isLoading}
-        draggable={isDraggable}
-        runQuery={runQuery}
-        results={results}
+        initialResults={results}
        onRunQuery={onRunQuery}
        onResults={onResults}
-        onDragStart={onDragStart}
        onUpdateChartConfig={onUpdateChartConfig}
+        draggable={isDraggable}
+        onDragStart={onDragStart}
      />
    </DebouncedComponent>
  )
@@ -10,16 +10,14 @@ import { ChartConfig } from 'components/interfaces/SQLEditor/UtilityPanel/ChartC
 import Results from 'components/interfaces/SQLEditor/UtilityPanel/Results'
 import { usePrimaryDatabase } from 'data/read-replicas/replicas-query'
 import { type QueryResponseError, useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
-import { type Parameter, parseParameters } from 'lib/sql-parameters'
-import type { Dashboards } from 'types'
-import { ChartContainer, ChartTooltipContent, cn, CodeBlock, SQL_ICON } from 'ui'
+
+import { Badge, ChartContainer, ChartTooltipContent, cn, CodeBlock, SQL_ICON } from 'ui'
 import ShimmeringLoader from 'ui-patterns/ShimmeringLoader'
 import { ButtonTooltip } from '../ButtonTooltip'
 import { CHART_COLORS } from '../Charts/Charts.constants'
 import SqlWarningAdmonition from '../SqlWarningAdmonition'
 import { BlockViewConfiguration } from './BlockViewConfiguration'
 import { EditQueryButton } from './EditQueryButton'
-import { ParametersPopover } from './ParametersPopover'
 import { getCumulativeResults } from './QueryBlock.utils'

 export const DEFAULT_CHART_CONFIG: ChartConfig = {
@@ -33,68 +31,21 @@ export const DEFAULT_CHART_CONFIG: ChartConfig = {
 }

 interface QueryBlockProps {
-  /** Applicable if SQL is a snippet that's already saved (Used in Reports) */
  id?: string
-  /** Title of the QueryBlock */
  label: string
-  /** SQL query to render/run in the QueryBlock */
  sql?: string
-  /** Configuration of the output chart based on the query result */
+  isWriteQuery?: boolean
  chartConfig?: ChartConfig
-  /** Not implemented yet: Will be the next part of ReportsV2 */
-  parameterValues?: Record<string, string>
-  /** Any other actions specific to the parent to be rendered in the header */
  actions?: ReactNode
-  /** Toggle visiblity of SQL query on render */
-  showSql?: boolean
-  /** Indicate if SQL query can be rendered as a chart */
-  isChart?: boolean
-  /** For Assistant as QueryBlock is rendered while streaming response */
-  isLoading?: boolean
-  /** Override to prevent running the SQL query provided */
-  runQuery?: boolean
-  /** Prevent updating of columns for X and Y axes in the chart view */
-  lockColumns?: boolean
-  /** Max height set to render results / charts (Defaults to 250) */
-  maxHeight?: number
-  /** Whether query block is draggable */
-  draggable?: boolean
-  /** Tooltip when hovering over the header of the block (Used in Assistant Panel) */
-  tooltip?: ReactNode
-  /** Optional: Any initial results to render as part of the query*/
-  results?: any[]
-  /** Opt to show run button if query is not read only */
-  showRunButtonIfNotReadOnly?: boolean
-  /** Not implemented yet: Will be the next part of ReportsV2 */
-  onSetParameter?: (params: Parameter[]) => void
-  /** Optional callback the SQL query is run */
-  onRunQuery?: (queryType: 'select' | 'mutation') => void
-  /** Optional callback on drag start */
-  onDragStart?: (e: DragEvent<Element>) => void
-  /** Optional: callback when the results are returned from running the SQL query*/
+  initialResults?: any[]
  onResults?: (results: any[]) => void
-  /** Optional: callback when an error occurs while running the SQL query*/
+  onRunQuery?: (queryType: 'select' | 'mutation') => void
  onError?: (errorText: string) => void
-  /** Force showing confirmation (write) bar before executing */
-  forceShowConfirmation?: boolean
-
-  // [Joshen] Params below are currently only used by ReportsV2 (Might revisit to see how to improve these)
-  /** Optional height set to render the SQL query (Used in Reports) */
-  queryHeight?: number
-  /** UI to render if there's a read-only error while running the query */
-  readOnlyErrorPlaceholder?: ReactNode
-  /** UI to render if there's no query results (Used in Reports) */
-  noResultPlaceholder?: ReactNode
-  /** To trigger a refresh of the query */
-  isRefreshing?: boolean
-  /** Optional callback whenever a chart configuration is updated (Used in Reports) */
-  onUpdateChartConfig?: ({
-    chart,
-    chartConfig,
-  }: {
-    chart?: Partial<Dashboards.Chart>
-    chartConfig: Partial<ChartConfig>
-  }) => 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
@@ -104,44 +55,34 @@ export const QueryBlock = ({
  label,
  sql,
  chartConfig = DEFAULT_CHART_CONFIG,
-  maxHeight = 250,
-  queryHeight,
-  parameterValues: extParameterValues,
  actions,
-  showSql: _showSql = false,
-  isChart = false,
-  isLoading = false,
-  runQuery = false,
-  lockColumns = false,
-  draggable = false,
-  isRefreshing = false,
-  noResultPlaceholder = null,
-  readOnlyErrorPlaceholder = null,
-  showRunButtonIfNotReadOnly = false,
-  tooltip,
-  results,
+  initialResults,
+  isWriteQuery = false,
  onRunQuery,
-  onSetParameter,
  onUpdateChartConfig,
-  onDragStart,
  onResults,
  onError,
-  forceShowConfirmation,
+  draggable = false,
+  onDragStart,
+  disabled = false,
+  isExternallyExecuting = false,
 }: QueryBlockProps) => {
  const { ref } = useParams()

  const [chartSettings, setChartSettings] = useState<ChartConfig>(chartConfig)
  const { xKey, yKey, view = 'table' } = chartSettings

-  const [showSql, setShowSql] = useState(_showSql)
+  // Maintain local results state, seeded from initial results prop
+  const [localResults, setLocalResults] = useState<any[] | undefined>(initialResults)
+
+  const [showSql, setShowSql] = useState(!localResults)
  const [readOnlyError, setReadOnlyError] = useState(false)
  const [queryError, setQueryError] = useState<QueryResponseError>()
-  const [queryResult, setQueryResult] = useState<any[] | undefined>(results)
  const [focusDataIndex, setFocusDataIndex] = useState<number>()

  const formattedQueryResult = useMemo(() => {
    // Make sure Y axis values are numbers
-    return queryResult?.map((row) => {
+    return localResults?.map((row) => {
      return Object.fromEntries(
        Object.entries(row).map(([key, value]) => {
          if (key === yKey) return [key, Number(value)]
@@ -149,18 +90,10 @@ export const QueryBlock = ({
        })
      )
    })
-  }, [queryResult, yKey])
+  }, [localResults, yKey])

-  const [parameterValues, setParameterValues] = useState<Record<string, string>>({})
  const [showWarning, setShowWarning] = useState<'hasWriteOperation' | 'hasUnknownFunctions'>()

-  const parameters = useMemo(() => {
-    if (!sql) return []
-    return parseParameters(sql)
-  }, [sql])
-  // [Joshen] This is for when we introduced the concept of parameters into our reports
-  // const combinedParameterValues = { ...extParameterValues, ...parameterValues }
-
  const { database: primaryDatabase } = usePrimaryDatabase({ projectRef: ref })
  const postgresConnectionString = primaryDatabase?.connectionString
  const readOnlyConnectionString = primaryDatabase?.connection_string_read_only
@@ -169,10 +102,10 @@ export const QueryBlock = ({
    ? getCumulativeResults({ rows: formattedQueryResult ?? [] }, chartSettings)
    : formattedQueryResult

-  const { mutate: execute, isLoading: isExecuting } = useExecuteSqlMutation({
+  const { mutate: execute, isLoading: isInternallyExecuting } = useExecuteSqlMutation({
    onSuccess: (data) => {
+      setLocalResults(data.result)
      onResults?.(data.result)
-      setQueryResult(data.result)

      setReadOnlyError(false)
      setQueryError(undefined)
@@ -183,7 +116,7 @@ export const QueryBlock = ({
      const notOwner = error.message.includes('must be owner')
      if (readOnlyTransaction || permissionDenied || notOwner) {
        setReadOnlyError(true)
-        if (showRunButtonIfNotReadOnly) setShowWarning('hasWriteOperation')
+        setShowWarning('hasWriteOperation')
      } else {
        setQueryError(error)
      }
@@ -199,16 +132,17 @@ export const QueryBlock = ({
  }
  const xKeyDateFormat = getDateFormat(xKey)

+  // Combined loading state from internal execution or external execution
+  const isExecuting = isInternallyExecuting || isExternallyExecuting
+
  const handleExecute = useCallback(() => {
    if (!sql || isExecuting) return

-    if (readOnlyError) {
+    if (readOnlyError || isWriteQuery) {
      return setShowWarning('hasWriteOperation')
    }

    try {
-      // [Joshen] This is for when we introduced the concept of parameters into our reports
-      // const processedSql = processParameterizedSql(sql, combinedParameterValues)
      execute({
        projectRef: ref,
        connectionString: readOnlyConnectionString,
@@ -217,116 +151,72 @@ export const QueryBlock = ({
    } catch (error: any) {
      toast.error(`Failed to execute query: ${error.message}`)
    }
-  }, [sql, isExecuting, readOnlyError, execute, ref, readOnlyConnectionString])
-
-  useEffect(() => {
-    if (forceShowConfirmation) setShowWarning('hasWriteOperation')
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [forceShowConfirmation])
+  }, [sql, isExecuting, readOnlyError, isWriteQuery, execute, ref, readOnlyConnectionString])

  useEffect(() => {
    setChartSettings(chartConfig)
  }, [chartConfig])

-  // Run once on mount to parse parameters and notify parent
+  // Sync local results when initialResults changes (from external execution)
  useEffect(() => {
-    if (!!sql && onSetParameter) {
-      const params = parseParameters(sql)
-      onSetParameter(params)
+    if (initialResults !== undefined) {
+      setLocalResults(initialResults)
    }
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [sql])
+  }, [initialResults])

-  useEffect(() => {
-    if (!!sql && runQuery && !!readOnlyConnectionString && !readOnlyError && !isExecuting) {
-      handleExecute()
-    }
-  }, [
-    sql,
-    isLoading,
-    runQuery,
-    readOnlyConnectionString,
-    isExecuting,
-    readOnlyError,
-    handleExecute,
-  ])
-
-  useEffect(() => {
-    if (isRefreshing) handleExecute()
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [isRefreshing])
+  // Execution is only user-driven within this component

  return (
    <ReportBlockContainer
      draggable={draggable}
      showDragHandle={draggable}
-      tooltip={tooltip}
-      loading={isExecuting}
      onDragStart={(e: DragEvent<Element>) => onDragStart?.(e)}
-      icon={
-        <SQL_ICON
-          size={18}
-          strokeWidth={1.5}
-          className={cn(
-            'transition-colors fill-foreground-muted group-aria-selected:fill-foreground',
-            'w-5 h-5 shrink-0 grow-0 -ml-0.5'
-          )}
-        />
+      loading={isExecuting}
+      label={
+        <div className="flex items-center gap-2">
+          <span>{label}</span>
+          {isWriteQuery && <Badge variant="warning">Write</Badge>}
+        </div>
      }
-      label={label}
      actions={
-        <>
-          <ButtonTooltip
-            type="text"
-            size="tiny"
-            className="w-7 h-7"
-            icon={<Code size={14} strokeWidth={1.5} />}
-            onClick={() => setShowSql(!showSql)}
-            tooltip={{
-              content: { side: 'bottom', text: showSql ? 'Hide query' : 'Show query' },
-            }}
-          />
+        disabled ? null : (
+          <>
+            <ButtonTooltip
+              type="text"
+              size="tiny"
+              className="w-7 h-7"
+              icon={<Code size={14} strokeWidth={1.5} />}
+              onClick={() => setShowSql(!showSql)}
+              tooltip={{
+                content: { side: 'bottom', text: showSql ? 'Hide query' : 'Show query' },
+              }}
+            />
+            {localResults && (
+              <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 })
+                }}
+                updateChartConfig={(config) => {
+                  if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: config })
+                  setChartSettings(config)
+                }}
+              />
+            )}

-          {queryResult && (
-            <>
-              {/* [Joshen ReportsV2] Won't see this just yet as this is intended for Reports V2 */}
-              {parameters.length > 0 && (
-                <ParametersPopover
-                  parameters={parameters}
-                  parameterValues={parameterValues}
-                  onSubmit={setParameterValues}
-                />
-              )}
-              {isChart && (
-                <BlockViewConfiguration
-                  view={view}
-                  isChart={isChart}
-                  lockColumns={lockColumns}
-                  chartConfig={chartSettings}
-                  columns={Object.keys(queryResult[0] || {})}
-                  changeView={(view) => {
-                    if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: { view } })
-                    setChartSettings({ ...chartSettings, view })
-                  }}
-                  updateChartConfig={(config) => {
-                    if (onUpdateChartConfig) onUpdateChartConfig({ chartConfig: config })
-                    setChartSettings(config)
-                  }}
-                />
-              )}
-            </>
-          )}
-
-          <EditQueryButton id={id} title={label} sql={sql} />
-
-          {(showRunButtonIfNotReadOnly || !readOnlyError) && (
+            <EditQueryButton id={id} title={label} sql={sql} />
            <ButtonTooltip
              type="text"
              size="tiny"
              className="w-7 h-7"
              icon={<Play size={14} strokeWidth={1.5} />}
-              loading={isExecuting || isLoading}
-              disabled={isLoading}
+              loading={isExecuting}
+              disabled={isExecuting}
              onClick={() => {
                handleExecute()
                if (!!sql) onRunQuery?.('select')
@@ -343,10 +233,10 @@ export const QueryBlock = ({
                },
              }}
            />
-          )}

-          {actions}
-        </>
+            {actions}
+          </>
+        )
      }
    >
      {!!showWarning && (
@@ -383,9 +273,8 @@ export const QueryBlock = ({
      {showSql && (
        <div
          className={cn('shrink-0 w-full max-h-96 overflow-y-auto', {
-            'border-b': queryResult !== undefined,
+            'border-b': localResults !== undefined,
          })}
-          style={{ height: !!queryHeight ? `${queryHeight}px` : undefined }}
        >
          <CodeBlock
            hideLineNumbers
@@ -400,15 +289,15 @@ export const QueryBlock = ({
        </div>
      )}

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

-      {view === 'chart' && queryResult !== undefined ? (
+      {view === 'chart' && localResults !== undefined ? (
        <>
-          {(queryResult ?? []).length === 0 ? (
+          {(localResults ?? []).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>
@@ -420,10 +309,7 @@ export const QueryBlock = ({
            <div className="flex-1 w-full">
              <ChartContainer
                className="aspect-auto px-3 py-2"
-                style={{
-                  height: maxHeight ? `${maxHeight}px` : undefined,
-                  minHeight: maxHeight ? `${maxHeight}px` : undefined,
-                }}
+                style={{ height: '250px', minHeight: '250px' }}
              >
                <BarChart
                  accessibilityLayer
@@ -469,26 +355,16 @@ export const QueryBlock = ({
      ) : (
        <>
          {!isExecuting && !!queryError ? (
-            <div
-              className={cn('flex-1 w-full overflow-auto relative border-t px-3.5 py-2')}
-              style={{ maxHeight: maxHeight ? `${maxHeight}px` : undefined }}
-            >
+            <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>
            </div>
-          ) : queryResult ? (
-            <div
-              className={cn('flex-1 w-full overflow-auto relative')}
-              style={{ maxHeight: maxHeight ? `${maxHeight}px` : undefined }}
-            >
-              <Results rows={queryResult} />
-            </div>
-          ) : !isExecuting ? (
-            readOnlyError ? (
-              readOnlyErrorPlaceholder
-            ) : (
-              noResultPlaceholder
+          ) : (
+            localResults && (
+              <div className={cn('flex-1 w-full overflow-auto relative max-h-64')}>
+                <Results rows={localResults} />
+              </div>
            )
-          ) : null}
+          )}
        </>
      )}
    </ReportBlockContainer>
@@ -280,41 +280,41 @@ Developer: # Role and Objective
 - Be aware that tool access may be restricted depending on the user's organization settings.
 - Do not try to bypass tool restrictions by executing SQL e.g. writing a query to retrieve database schema information. Instead, explain to the user you do not have permissions to use the tools you need to execute the task

-# Output Format
+## Output Format
 - Always integrate findings from the tools seamlessly into your responses for better accuracy and context.

-# Searching Docs
+##  Searching Docs
 - Use \`search_docs\` to search the Supabase documentation for relevant information when the question is about Supabase features or complex database operations
 `

 export const CHAT_PROMPT = `
-Developer: # Response Style
- Be direct and concise. Provide only essential information.
- Use lists to present information; do not use tables for formatting.
- Minimize use of emojis.
+# Response Style
+- Be professional, direct and concise. Provide only essential information.
+- Before context gathering or tool usage, summarise the user's request and your plan of action in a single paragraph
+- Do not repeat yourself or your plan after context gathering

 # Response Format
-## Markdown
- Follow the CommonMark specification.
- Use a logical heading hierarchy (H1–H4), maintaining order without skipping levels.
+## Use Markdown
+- *CRITICAL*: Response must be in markdown format.
+- Always use markdown blocks **where semantically correct** (e.g., \`inline code\`, \`\`\`code fences\`\`\`, headings, lists, tables).
+- Make use of markdown headings to structure your response where appropriate. e.g. WRONG "Section heading: ..." WRITE "## Section heading"
+- Shorter responses do not need headings
 - Use bold text exclusively to emphasize key information.
 - Do not use tables for displaying information under any circumstances.
+- Minimize use of emojis.

 # Chat Naming
 - At the start of each conversation, if the chat has not yet been named, invoke \`rename_chat\` with a descriptive 2–4 word name. Examples: "User Authentication Setup", "Sales Data Analysis", "Product Table Creation".

-## Task Workflow
- Always start the conversation with a concise checklist of sub-tasks you will perform before generating outputs or calling tools. Keep the checklist conceptual, not implementation-level.
- No need to repeat the checklist later in the conversation
-
 # SQL Execution and Display
 - Be confident: assume the user is the project owner. You do not need to show code before execution.
 - To actually run SQL, directly call the \`execute_sql\` tool with the \`sql\` string. The client will request user confirmation and then return results.
 - If executing SQL returns an error, explain the error concisely and try again with the correct SQL.
- The user may skip executing the query in which case you should offer alternative options or actions to take
+- The user may skip executing the query in which case you should acknowledge the skip and offer alternative options or actions to take
 - If the user asks you to write a query, or if you want to show example SQL without executing, render it in a markdown code block (e.g.: \`\`\`sql). Do this only when the user asks to see the code or for illustrative examples.
 - If multiple queries are needed, call \`execute_sql\` separately for each and validate results in 1–2 lines. Use separate code blocks only for non-executed examples.
 - After executing queries, summarize the outcome and confirm next actions or self-correct as needed.
+- You do not need to repeat the SQL query results as the client will display them to the user as part of the execute_sql tool call.

 # Edge Functions
 - Be confident: assume the user is the project owner. You do not need to show code before deployment.
@@ -3,6 +3,8 @@ import { createSupabaseMCPClient } from '../supabase-mcp'
 import { filterToolsByOptInLevel, toolSetValidationSchema } from '../tool-filter'
 import type { ToolSet } from 'ai'

+const UI_EXECUTED_TOOLS = ['execute_sql', 'deploy_edge_function']
+
 export const getMcpTools = async ({
  accessToken,
  projectRef,
@@ -24,8 +26,9 @@ export const getMcpTools = async ({

  // Remove UI-executed tools handled locally
  const filteredMcpTools: ToolSet = { ...allowedMcpTools }
-  delete (filteredMcpTools as any)['execute_sql']
-  delete (filteredMcpTools as any)['deploy_edge_function']
+  UI_EXECUTED_TOOLS.forEach((toolName) => {
+    delete filteredMcpTools[toolName]
+  })

  // Validate that only known tools are provided
  const validation = toolSetValidationSchema.safeParse(filteredMcpTools)
@@ -6,8 +6,13 @@ export const getRenderingTools = () => ({
    description: 'Asks the user to execute a SQL statement and return the results',
    inputSchema: z.object({
      sql: z.string().describe('The SQL statement to execute.'),
+      label: z.string().describe('A short 2-4 word label for the SQL statement.'),
+      isWriteQuery: z
+        .boolean()
+        .describe(
+          'Whether the SQL statement performs a write operation of any kind instead of a read operation'
+        ),
    }),
-    // No execute – runs client-side in the UI
  }),
  deploy_edge_function: tool({
    description:
@@ -16,7 +21,6 @@ export const getRenderingTools = () => ({
      name: z.string().describe('The URL-friendly name/slug of the Edge Function.'),
      code: z.string().describe('The TypeScript code for the Edge Function.'),
    }),
-    // No execute – runs client-side in the UI
  }),
  rename_chat: tool({
    description: `Rename the current chat session when the current chat name doesn't describe the conversation topic.`,
@@ -72,26 +72,6 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {

  const { messages: rawMessages, projectRef, connectionString, orgSlug, chatName } = data

-  // Server-side safety: limit to last 7 messages and remove `results` property to prevent accidental leakage.
-  // Results property is used to cache results client-side after queries are run
-  // Tool results will still be included in history sent to model
-  const messages = (rawMessages || []).slice(-7).map((msg: any) => {
-    if (msg && msg.role === 'assistant' && 'results' in msg) {
-      const cleanedMsg = { ...msg }
-      delete cleanedMsg.results
-      return cleanedMsg
-    }
-    // [Joshen] Am also filtering out any tool calls which state is "input-streaming"
-    // this happens when a user stops the assistant response while the tool is being called
-    if (msg && msg.role === 'assistant' && msg.parts) {
-      const cleanedParts = msg.parts.filter((part: any) => {
-        return !(part.type.startsWith('tool-') && part.state === 'input-streaming')
-      })
-      return { ...msg, parts: cleanedParts }
-    }
-    return msg
-  })
-
  let aiOptInLevel: AiOptInLevel = 'disabled'
  let isLimited = false

@@ -117,6 +97,39 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) {
    }
  }

+  // Server-side safety: limit to last 7 messages and remove `results` property to prevent accidental leakage.
+  // Results property is used to cache results client-side after queries are run
+  // Tool results will still be included in history sent to model
+  const messages = (rawMessages || []).slice(-7).map((msg: any) => {
+    if (msg && msg.role === 'assistant' && 'results' in msg) {
+      const cleanedMsg = { ...msg }
+      delete cleanedMsg.results
+      return cleanedMsg
+    }
+    // Filter out tool calls with UI-specific states that the AI SDK doesn't understand
+    // These states are used for UI rendering but shouldn't be sent to the model
+    if (msg && msg.role === 'assistant' && msg.parts) {
+      const cleanedParts = msg.parts
+        .filter((part: any) => {
+          // Filter out tool calls with invalid states for AI SDK
+          if (part.type.startsWith('tool-')) {
+            const invalidStates = ['input-streaming', 'input-available', 'output-error']
+            return !invalidStates.includes(part.state)
+          }
+          return true
+        })
+        .map((part: any) => {
+          // If opt in level is schema_and_log_and_data, replace part output with permission message
+          if (aiOptInLevel !== 'schema_and_log_and_data' && part.output) {
+            return { ...part, output: 'Result not shared due to permission level' }
+          }
+          return part
+        })
+      return { ...msg, parts: cleanedParts }
+    }
+    return msg
+  })
+
  const {
    model,
    error: modelError,
@@ -104,16 +104,34 @@ async function clearStorage(): Promise<void> {
  }
 }

+// Helper function to sanitize objects to ensure they're cloneable
+function sanitizeForCloning(obj: any): any {
+  if (obj === null || obj === undefined) return obj
+  if (typeof obj !== 'object') return obj
+
+  // Convert to plain object to remove any proxy wrappers
+  return JSON.parse(JSON.stringify(obj))
+}
+
 // Helper function to load state from IndexedDB
 async function loadFromIndexedDB(projectRef: string): Promise<StoredAiAssistantState | null> {
  try {
    const persistedState = await getAiState(projectRef)
    if (persistedState) {
-      // Revive dates
+      // Revive dates and sanitize message data
      Object.values(persistedState.chats).forEach((chat: ChatSession) => {
        if (chat && typeof chat === 'object') {
          chat.createdAt = new Date(chat.createdAt)
          chat.updatedAt = new Date(chat.updatedAt)
+
+          // Sanitize message parts to remove proxy objects
+          if (chat.messages) {
+            chat.messages.forEach((message: any) => {
+              if (message.parts) {
+                message.parts = message.parts.map((part: any) => sanitizeForCloning(part))
+              }
+            })
+          }
        }
      })
      return persistedState
@@ -337,26 +355,15 @@ export const createAiAssistantState = (): AiAssistantState => {
      }
    },

-    updateMessage: ({
-      id,
-      resultId,
-      results,
-    }: {
-      id: string
-      resultId?: string
-      results: any[]
-    }) => {
+    updateMessage: (updatedMessage: MessageType) => {
      const chat = state.activeChat
-      if (!chat || !resultId) return
-
-      const messageIndex = chat.messages.findIndex((msg) => msg.id === id)
+      if (!chat) return

+      const messageIndex = chat.messages.findIndex((msg) => msg.id === updatedMessage.id)
      if (messageIndex !== -1) {
-        const msg = chat.messages[messageIndex]
-        if (!msg.results) {
-          msg.results = {}
-        }
-        msg.results[resultId] = results
+        // Replace entire message instance to capture latest tool parts/results
+        chat.messages[messageIndex] = updatedMessage as AssistantMessageType
+        chat.updatedAt = new Date()
      }
    },

@@ -435,7 +442,7 @@ export type AiAssistantState = AiAssistantData & {
  clearMessages: () => void
  deleteMessagesAfter: (id: string, options?: { includeSelf?: boolean }) => void
  saveMessage: (message: MessageType | MessageType[]) => void
-  updateMessage: (args: { id: string; resultId?: string; results: any[] }) => void
+  updateMessage: (message: MessageType) => void
  setSqlSnippets: (snippets: SqlSnippet[]) => void
  clearSqlSnippets: () => void
  getCachedSQLResults: (args: { messageId: string; snippetId?: string }) => any[] | undefined