[pull] master from supabase:master

apps/studio/components/interfaces/Explorer/QueryEditor/QueryResultChart.tsx

@@ -81,45 +81,43 @@ export const QueryResultChart = ({ chart, result }: QueryResultChartProps) => {
   }
 
   return (
-    <Chart>
-      <ChartCard className="rounded-none border-0">
-        <ChartContent>
-          <div className="h-40">
-            {type === 'bar' && (
-              <ChartBar
-                isFullHeight
-                xKey={x_column}
-                dataKey={y_columns[0]}
-                dataKeys={y_columns}
-                config={chartConfig}
-                showXAxis={show_labels}
-                showYAxis={show_labels}
-                data={resultToRender}
-                YAxisProps={{
-                  scale: effectiveScale === 'log' ? 'log' : 'auto',
-                  domain: effectiveScale === 'log' ? [1, 'auto'] : undefined,
-                  tickFormatter: effectiveScale === 'log' ? formatLogTick : undefined,
-                }}
-              />
-            )}
-            {type === 'line' && (
-              <ChartLine
-                isFullHeight
-                xKey={x_column}
-                dataKey={y_columns[0]}
-                dataKeys={y_columns}
-                config={chartConfig}
-                showXAxis={show_labels}
-                showYAxis={show_labels}
-                data={resultToRender}
-                YAxisProps={{
-                  scale: effectiveScale === 'log' ? 'log' : 'auto',
-                  domain: effectiveScale === 'log' ? [1, 'auto'] : undefined,
-                  tickFormatter: effectiveScale === 'log' ? formatLogTick : undefined,
-                }}
-              />
-            )}
-          </div>
+    <Chart className="flex h-full min-h-0 flex-col">
+      <ChartCard className="flex h-full min-h-0 flex-col rounded-none border-0">
+        <ChartContent className="flex h-full min-h-0 flex-1 flex-col">
+          {type === 'bar' && (
+            <ChartBar
+              isFullHeight
+              xKey={x_column}
+              dataKey={y_columns[0]}
+              dataKeys={y_columns}
+              config={chartConfig}
+              showXAxis={show_labels}
+              showYAxis={show_labels}
+              data={resultToRender}
+              YAxisProps={{
+                scale: effectiveScale === 'log' ? 'log' : 'auto',
+                domain: effectiveScale === 'log' ? [1, 'auto'] : undefined,
+                tickFormatter: effectiveScale === 'log' ? formatLogTick : undefined,
+              }}
+            />
+          )}
+          {type === 'line' && (
+            <ChartLine
+              isFullHeight
+              xKey={x_column}
+              dataKey={y_columns[0]}
+              dataKeys={y_columns}
+              config={chartConfig}
+              showXAxis={show_labels}
+              showYAxis={show_labels}
+              data={resultToRender}
+              YAxisProps={{
+                scale: effectiveScale === 'log' ? 'log' : 'auto',
+                domain: effectiveScale === 'log' ? [1, 'auto'] : undefined,
+                tickFormatter: effectiveScale === 'log' ? formatLogTick : undefined,
+              }}
+            />
+          )}
         </ChartContent>
       </ChartCard>
     </Chart>

apps/studio/components/interfaces/Explorer/QueryEditor/index.tsx

@@ -80,13 +80,17 @@ export type QueryEditorProps = {
   roleImpersonationState?: RoleImpersonationController
   display?: QueryDisplay
   toolbarActions?: ReactNode
+  className?: string
+  /** When true, toolbar and editor run actions are disabled. */
+  isRunDisabled?: boolean
   onTitleChange: (title: string) => void
   onSqlChange: (sql: string) => void
   onSqlCommit?: (sql: string) => void
   onSourceChange?: (source: QuerySourceBinding) => void
   onResultChange: (result: QueryResult) => void
   onRowLimitChange?: (val: number) => void
   onDisplayChange?: (display: QueryDisplay) => void
+  onRun?: () => void
 }
 
 export type QueryEditorHandle = {
@@ -108,13 +112,16 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
     roleImpersonationState,
     display,
     toolbarActions,
+    className,
+    isRunDisabled = false,
     onTitleChange,
     onSqlChange,
     onSqlCommit,
     onSourceChange,
     onResultChange,
     onRowLimitChange,
     onDisplayChange,
+    onRun,
   }: QueryEditorProps,
   ref
 ) {
@@ -166,8 +173,9 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
    * Postgres SQL cannot reach the analytics wire or vice versa.
    */
   const handleRunQuery = async (rawSql: string = sql) => {
-    if (!project || isBusy || rewriteProposal || rawSql.trim().length === 0) return
+    if (!project || isBusy || rewriteProposal || isRunDisabled || rawSql.trim().length === 0) return
 
+    onRun?.()
     onSqlCommit?.(rawSql)
 
     if (query._tag === 'logs') {
@@ -227,7 +235,7 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
   const Shell = variant === 'viewport' ? ExplorerQueryViewport : ExplorerQuery
 
   return (
-    <Shell className={variant === 'embedded' ? 'mx-auto max-w-4xl' : undefined}>
+    <Shell className={cn(variant === 'embedded' && 'mx-auto max-w-4xl', className)}>
       <ExplorerToolbar>
         <ExplorerToolbarIcon>
           <CodeSquare size={14} />
@@ -268,7 +276,11 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
             icon={<Play />}
             tooltip="Run query"
             disabled={
-              isLoadingProject || isExecuting || rewriteProposal !== null || sql.trim().length === 0
+              isLoadingProject ||
+              isExecuting ||
+              rewriteProposal !== null ||
+              isRunDisabled ||
+              sql.trim().length === 0
             }
             onClick={() => handleRunQuery()}
           >
@@ -296,7 +308,7 @@ export const QueryEditor = forwardRef<QueryEditorHandle, QueryEditorProps>(funct
               placeholder="select * from your_table limit 100;"
               placeholderClassName="top-[13px]"
               className={variant === 'embedded' ? 'h-44' : undefined}
-              actions={{ runQuery: { enabled: true, callback: handleRunQuery } }}
+              actions={{ runQuery: { enabled: !isRunDisabled, callback: handleRunQuery } }}
               options={{ minimap: { enabled: false }, padding: { top: 8 } }}
               onInputChange={(value) => onSqlChange(value ?? '')}
               onMount={(editor) => {

apps/studio/components/interfaces/Reports/ReportBlock/ReportBlockContainer.tsx

@@ -11,6 +11,7 @@ interface ReportBlockContainerProps {
   draggable?: boolean
   showDragHandle?: boolean
   tooltip?: ReactNode
+  className?: string
   onDragStart?: (e: DragEvent) => void
 }
 
@@ -23,6 +24,7 @@ export const ReportBlockContainer = ({
   draggable = false,
   showDragHandle = false,
   tooltip,
+  className,
   onDragStart,
   children,
 }: PropsWithChildren<ReportBlockContainerProps>) => {
@@ -35,7 +37,10 @@ export const ReportBlockContainer = ({
       draggable={draggable}
       unselectable={draggable ? 'on' : undefined}
       onDragStart={onDragStart}
-      className="h-full flex flex-col overflow-hidden bg-surface-100 border-overlay relative rounded-sm border shadow-xs"
+      className={cn(
+        'h-full flex flex-col overflow-hidden bg-surface-100 border-overlay relative rounded-sm border shadow-xs',
+        className
+      )}
     >
       <Tooltip>
         <TooltipTrigger asChild>

apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.test.ts

@@ -127,6 +127,18 @@ describe('AIAssistant.utils.ts:hasPendingToolApproval', () => {
 
     expect(hasPendingToolApproval(messages)).toBe(true)
   })
+
+  test('Should ignore automatic approvals', () => {
+    const messages = createMessageWithPart({
+      type: 'tool-execute_sql',
+      toolCallId: 'call-1',
+      state: 'approval-requested',
+      input: { sql: 'select 1', label: 'Test query' },
+      approval: { id: 'approval-1', isAutomatic: true },
+    } as UIMessage['parts'][number])
+
+    expect(hasPendingToolApproval(messages)).toBe(false)
+  })
 })
 
 describe('AIAssistant.utils.ts:resolvePendingToolApprovalsAsDenied', () => {

apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts

@@ -15,6 +15,7 @@ import { databaseKeys } from '@/data/database/keys'
 import { enumeratedTypesKeys } from '@/data/enumerated-types/keys'
 import { handleError } from '@/data/fetchers'
 import { tableKeys } from '@/data/tables/keys'
+import { isManualApprovalRequested } from '@/lib/ai/message-utils'
 import { tryParseJson } from '@/lib/helpers'
 import type { SqlSnippet } from '@/state/ai-assistant-state'
 import { ResponseError } from '@/types'
@@ -85,7 +86,7 @@ export const hasPendingToolApproval = (messages: Pick<UIMessage, 'role' | 'parts
   return messages.some((message) => {
     if (message.role !== 'assistant') return false
 
-    return message.parts?.some((part) => isToolUIPart(part) && part.state === 'approval-requested')
+    return message.parts?.some((part) => isManualApprovalRequested(part))
   })
 }
 
@@ -183,7 +184,7 @@ export const getSnippetContent = (snippet: SqlSnippet): string =>
  * against the `logs` table — a single message can carry both dialects.
  *
  * It also keeps the two apart in the rendered message: MessageMarkdown treats a `sql`
- * fence as runnable Postgres (`DisplayBlockRenderer`, branded with `untrustedSql`),
+ * fence as runnable Postgres (`AssistantQueryCell`, branded with `untrustedSql`),
  * which a ClickHouse query must never be offered as.
  */
 function getSnippetFenceLanguage(snippet: SqlSnippet): 'sql' | 'clickhouse' {

apps/studio/components/ui/AIAssistantPanel/AssistantChat.tsx

@@ -194,7 +194,7 @@ export const AssistantChat = ({
     addToolApprovalResponse,
     stop,
     regenerate,
-  } = useChat({
+  } = useChat<MessageType>({
     id: chatId,
     ...(chatInstance ? { chat: chatInstance } : {}),
     sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,

apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.tsx

@@ -0,0 +1,143 @@
+import { useRef, useState } from 'react'
+
+import { identifyQueryType } from './AIAssistant.utils'
+import {
+  changeAssistantQuerySource,
+  createAssistantQueryModel,
+  DEFAULT_ASSISTANT_QUERY_TITLE,
+  getAssistantQueryDisplay,
+  setAssistantQuerySql,
+  toAssistantQueryResult,
+} from './AssistantQueryCell.utils'
+import { Confirm } from './Confirm'
+import { type ConfirmFooterApprovalState } from './Confirm.utils'
+import { QueryEditor } from '@/components/interfaces/Explorer/QueryEditor'
+import { type QueryDisplay, type QueryResult } from '@/components/interfaces/Explorer/types'
+import { type QuerySourceBinding } from '@/data/query-sources/query-source-registry'
+import { useTrack } from '@/lib/telemetry/track'
+import { useLocalRoleImpersonationState } from '@/state/role-impersonation-state'
+
+interface AssistantQueryCellProps {
+  id: string
+  sql: string
+  title?: string
+  initialRows?: unknown
+  view?: 'table' | 'chart'
+  xAxis?: string
+  yAxis?: string
+  /** Follow incoming SQL while the assistant is still streaming the query text. */
+  isStreaming?: boolean
+  confirmState?: ConfirmFooterApprovalState
+  onApprove?: () => void
+  onDeny?: () => void
+}
+
+/** Assistant adapter around the shared QueryEditor. Local state only — nothing is persisted. */
+export const AssistantQueryCell = ({
+  id,
+  sql: initialSql,
+  title: initialTitle,
+  initialRows,
+  view,
+  xAxis,
+  yAxis,
+  isStreaming = false,
+  confirmState,
+  onApprove,
+  onDeny,
+}: AssistantQueryCellProps) => {
+  const track = useTrack()
+  const roleImpersonationState = useLocalRoleImpersonationState()
+
+  const [title, setTitle] = useState(initialTitle?.trim() || DEFAULT_ASSISTANT_QUERY_TITLE)
+  const [query, setQuery] = useState(() => createAssistantQueryModel(initialSql))
+  const [result, setResult] = useState<QueryResult | undefined>(() =>
+    toAssistantQueryResult(initialRows)
+  )
+  const [display, setDisplay] = useState<QueryDisplay>(() =>
+    getAssistantQueryDisplay({ view, xAxis, yAxis })
+  )
+
+  const prevId = useRef(id)
+  const prevSql = useRef(initialSql)
+  const prevRows = useRef(initialRows)
+
+  if (prevId.current !== id) {
+    prevId.current = id
+    prevSql.current = initialSql
+    prevRows.current = initialRows
+    setTitle(initialTitle?.trim() || DEFAULT_ASSISTANT_QUERY_TITLE)
+    setQuery(createAssistantQueryModel(initialSql))
+    setResult(toAssistantQueryResult(initialRows))
+    setDisplay(getAssistantQueryDisplay({ view, xAxis, yAxis }))
+  }
+
+  if (prevSql.current !== initialSql) {
+    prevSql.current = initialSql
+    if (isStreaming) {
+      setQuery((current) => setAssistantQuerySql(current, initialSql))
+    }
+  }
+
+  if (prevRows.current !== initialRows) {
+    prevRows.current = initialRows
+    setResult(toAssistantQueryResult(initialRows))
+  }
+
+  const handleTitleChange = (value: string) => {
+    const nextTitle = value.trim()
+    if (!nextTitle) return
+    setTitle(nextTitle)
+  }
+
+  const handleSourceChange = (source: QuerySourceBinding) => {
+    const isBackendChange = source._tag !== query._tag
+    if (isBackendChange) setResult(undefined)
+    setQuery((current) => changeAssistantQuerySource(current, source))
+  }
+
+  const handleRun = () => {
+    const sql = query.uncheckedSql
+    const mutationType = identifyQueryType(sql)
+    track('assistant_suggestion_run_query_clicked', {
+      queryType: mutationType ? 'mutation' : 'select',
+      ...(mutationType ? { mutationType } : {}),
+    })
+  }
+
+  const isConfirming = confirmState !== undefined
+
+  return (
+    <Confirm
+      fill
+      className="h-96"
+      state={confirmState}
+      message="Assistant wants to run this query"
+      cancelLabel="Skip"
+      confirmLabel="Run query"
+      confirmLabelLoading="Running..."
+      onCancel={onDeny}
+      onConfirm={onApprove}
+    >
+      <QueryEditor
+        id={id}
+        variant="viewport"
+        title={title}
+        query={query}
+        result={result}
+        roleImpersonationState={roleImpersonationState}
+        display={display}
+        isRunDisabled={isConfirming}
+        onTitleChange={handleTitleChange}
+        onSqlChange={(sql) => setQuery((current) => setAssistantQuerySql(current, sql))}
+        onSourceChange={handleSourceChange}
+        onResultChange={setResult}
+        onRowLimitChange={(rowLimit) =>
+          setQuery((current) => (current._tag === 'database' ? { ...current, rowLimit } : current))
+        }
+        onDisplayChange={setDisplay}
+        onRun={handleRun}
+      />
+    </Confirm>
+  )
+}