feat(studio): debug query errors in the current assistant thread and fix prose alignment

Co-authored-by: Cursor <cursoragent@cursor.com>

apps/studio/components/interfaces/Explorer/QueryEditor/QueryResultError.test.tsx

@@ -105,7 +105,7 @@ describe('QueryResultError', () => {
     )
   })
 
-  it('calls onDebug with the prompt instead of opening a new chat when provided', () => {
+  it('hands the debug prompt to onDebug instead of opening a new chat', () => {
     const onDebug = vi.fn()
 
     customRender(
@@ -120,6 +120,7 @@ describe('QueryResultError', () => {
     fireEvent.click(screen.getByRole('button', { name: 'Debug with Assistant' }))
 
     expect(onDebug).toHaveBeenCalledWith(expect.stringContaining('select * from foo;'))
+    expect(onDebug.mock.calls[0][0]).toContain('relation "foo" does not exist')
     expect(mocks.createChat).not.toHaveBeenCalled()
   })
 

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

@@ -29,9 +29,7 @@ export const QueryResultError = ({
   autoLimit?: QueryResult['autoLimit']
   sql?: string
   source?: SqlSnippetSource
-  /** Overrides the default "open a new debug chat" behavior — used when this query block
-   * is already rendered inside an open assistant conversation, so debugging should write
-   * into that conversation's composer instead of abandoning it for a new chat. */
+  /** Receives the "Debug with Assistant" prompt. Defaults to opening a new assistant chat. */
   onDebug?: (prompt: string) => void
 }) => {
   const { ref } = useParams()
@@ -65,10 +63,11 @@ export const QueryResultError = ({
     [canDebug, sql, error.message, source]
   )
 
-  const handleDebug = () =>
-    onDebug
-      ? onDebug(buildDebugPrompt())
-      : createChat({ name: 'Debug SQL snippet', initialMessage: buildDebugPrompt() })
+  const handleDebug = () => {
+    const prompt = buildDebugPrompt()
+    if (onDebug) return onDebug(prompt)
+    createChat({ name: 'Debug SQL snippet', initialMessage: prompt })
+  }
 
   const isTimeout =
     error.message?.includes('canceling statement due to statement timeout') ||

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

@@ -141,6 +141,7 @@ type QueryEditorProps = {
   onRowLimitChange?: (val: number) => void
   onDisplayChange?: (display: QueryDisplay) => void
   onRun?: () => void
+  /** Receives the "Debug with Assistant" prompt on error. Defaults to opening a new assistant chat. */
   onDebug?: (prompt: string) => void
 }
 

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

@@ -8,6 +8,7 @@ import { Eraser, Pencil, X } from 'lucide-react'
 import { Fragment, useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from 'react'
 import { Button, cn, KeyboardShortcut } from 'ui'
 import { Admonition } from 'ui-patterns/Admonition'
+import type { StickToBottomContext } from 'use-stick-to-bottom'
 
 import { AlertError } from '../AlertError'
 import { ButtonTooltip } from '../ButtonTooltip'
@@ -121,6 +122,7 @@ export const AssistantChat = ({
 
   // Add a ref to store the last user message
   const lastUserMessageRef = useRef<MessageType | null>(null)
+  const conversationRef = useRef<StickToBottomContext>(null)
 
   const [value, setValue] = useState<string>(composerContext?.initialInput || '')
   const [editingMessageId, setEditingMessageId] = useState<string | null>(null)
@@ -388,6 +390,8 @@ export const AssistantChat = ({
       },
     })
     setValue('')
+    // Re-lock stick-to-bottom in case the user sent this from mid-thread (e.g. "Debug with Assistant").
+    conversationRef.current?.scrollToBottom()
 
     if (finalContent.includes('Help me to debug')) {
       track('assistant_debug_submitted', { chatId })
@@ -476,7 +480,7 @@ export const AssistantChat = ({
           aiOptInLevel,
         })}
         {hasMessages ? (
-          <Conversation className={cn('flex-1')}>
+          <Conversation className={cn('flex-1')} contextRef={conversationRef}>
             <ConversationContent className="w-full px-7 py-8 mb-10">
               {renderedMessages}
               <div className="w-full max-w-3xl mx-auto">
@@ -540,7 +544,7 @@ export const AssistantChat = ({
                 )}
 
                 <p className="text-center text-xs text-foreground-muted mt-6">
-                  The Assistant can make mistakes. Double check responses.
+                  Assistant can make mistakes. Double check responses.
                 </p>
               </div>
             </ConversationContent>

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

@@ -12,6 +12,7 @@ import {
 } from './AssistantQueryCell.utils'
 import { Confirm } from './Confirm'
 import { type ConfirmFooterApprovalState } from './Confirm.utils'
+import { useMessageActionsContext } from './Message.Context'
 import { QueryEditor } from '@/components/interfaces/Explorer/QueryEditor'
 import { type QueryDisplay, type QueryResult } from '@/components/interfaces/Explorer/types'
 import {
@@ -19,7 +20,6 @@ import {
   type QuerySourceTag,
 } from '@/data/query-sources/query-source-registry'
 import { useTrack } from '@/lib/telemetry/track'
-import { useAiAssistantState } from '@/state/ai-assistant-state'
 import { useLocalRoleImpersonationState } from '@/state/role-impersonation-state'
 
 interface AssistantQueryCellProps {
@@ -73,7 +73,8 @@ export const AssistantQueryCell = ({
 }: AssistantQueryCellProps) => {
   const track = useTrack()
   const roleImpersonationState = useLocalRoleImpersonationState()
-  const aiAssistantState = useAiAssistantState()
+  // "Debug with Assistant" continues the current thread rather than opening a new chat.
+  const { onSendMessage } = useMessageActionsContext()
 
   const fallbackTitle =
     initialTitle?.trim() ||
@@ -184,7 +185,7 @@ export const AssistantQueryCell = ({
         }
         onDisplayChange={handleDisplayChange}
         onRun={handleRun}
-        onDebug={aiAssistantState.setInitialInput}
+        onDebug={onSendMessage}
       />
     </Confirm>
   )

apps/studio/components/ui/AIAssistantPanel/Message.Context.tsx

@@ -29,6 +29,8 @@ export interface MessageActions {
   onBranch: (id: string) => void
   onCancelEdit: () => void
   onRate?: (id: string, rating: 'positive' | 'negative', reason?: string) => void
+  /** Sends a follow-up user message in the current chat. */
+  onSendMessage?: (text: string) => void
 }
 
 const MessageInfoContext = createContext<MessageInfo | null>(null)

apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx

@@ -36,6 +36,8 @@ function MessagePartText({ textPart }: { textPart: TextUIPart }) {
       readOnly={readOnly}
       className={cn(
         'max-w-none space-y-4 prose prose-sm prose-li:mt-1 [&>div]:my-4 prose-h1:text-xl prose-h1:mt-6 prose-h2:text-lg prose-h2:font-medium 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 wrap-break-word [&>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!',
+        // Prose stays in the narrow column; only blocks flagged wide (query cells) use the full width.
+        '[&>*]:mx-auto [&>*]:max-w-3xl [&>[data-wide]]:max-w-none',
         isUserMessage && 'text-foreground [&>p]:font-medium',
         state === 'editing' && 'animate-pulse'
       )}
@@ -318,8 +320,9 @@ const isWideMessagePart = (part: NonNullable<VercelMessage['parts']>[number]) =>
   part.type === 'tool-update_notebook' ||
   part.type === 'tool-run_notebook' ||
   (part.type === 'dynamic-tool' && part.toolName === 'query_logs') ||
-  // Unlabelled code fences resolve to SQL in MessageMarkdown, too.
-  (part.type === 'text' && /```(?:sql)?(?:\s|$)/i.test(part.text))
+  // Text spans the wide column so a SQL fence can render as a full-width query cell;
+  // MessagePartText re-centers its prose blocks at the narrow width.
+  part.type === 'text'
 
 const isCompactToolPart = (part: NonNullable<VercelMessage['parts']>[number]) =>
   part.type === 'reasoning' ||

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

@@ -111,6 +111,7 @@ interface MessageProps {
   onCancelEdit: () => void
   isLastMessage?: boolean
   onRate?: (id: string, rating: 'positive' | 'negative', reason?: string) => void
+  onSendMessage?: (text: string) => void
   rating?: 'positive' | 'negative' | null
 }
 
@@ -141,6 +142,7 @@ export function Message(props: MessageProps) {
     onBranch: props.onBranch,
     onCancelEdit: props.onCancelEdit,
     onRate: props.onRate,
+    onSendMessage: props.onSendMessage,
   }
 
   return (

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

@@ -218,13 +218,15 @@ export const MarkdownPre = ({
   const cleanContent = rawContent.replace(/(?:--|\/\/)\s*props:\s*\{[^}]+\}/, '').trim()
 
   const toolCallId = String(snippetId ?? id)
+  const isQueryCell = language === 'sql' && !readOnly
 
   if (!codeElement) {
     return <pre className="w-auto overflow-x-auto not-prose my-4">{children}</pre>
   }
 
   return (
-    <div className="w-auto overflow-x-hidden not-prose my-4 ">
+    // `data-wide` lets the parent text part give query cells the full column width.
+    <div className="w-auto overflow-x-hidden not-prose my-4" data-wide={isQueryCell || undefined}>
       {language === 'edge' ? (
         <EdgeFunctionBlock
           label={title}