feat(studio): render assistant notebook runs

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

@@ -18,6 +18,7 @@ import {
 import { MessageMarkdown } from './MessageMarkdown'
 import { MessagePartQueryLogs } from './MessagePartQueryLogs'
 import { NotebookProposalRenderer, type NotebookProposalMode } from './NotebookProposalRenderer'
+import { NotebookRunRenderer } from './NotebookRunRenderer'
 import { parseSupportRequestMessage, SupportRequestMessage } from './SupportRequestMessage'
 
 function MessagePartText({ textPart }: { textPart: TextUIPart }) {
@@ -261,6 +262,32 @@ function MessagePartNotebookProposal({
   )
 }
 
+function MessagePartNotebookRun({ toolPart }: { toolPart: ToolUIPart }) {
+  const { state, input: submittedInput, output } = toolPart
+  const input = state === 'output-error' ? (submittedInput ?? toolPart.rawInput) : submittedInput
+  const { addToolApprovalResponse } = useMessageActionsContext()
+
+  if (state === 'input-streaming')
+    return <ToolDisplayExecuteSqlLoading label="Preparing notebook..." />
+
+  const { confirmState, onApprove, onDeny } = getManualToolApprovalHandlers({
+    state,
+    approval: toolPart.approval,
+    addToolApprovalResponse,
+  })
+
+  return (
+    <NotebookRunRenderer
+      state={state}
+      input={input}
+      output={output}
+      confirmState={confirmState}
+      onApprove={onApprove}
+      onDeny={onDeny}
+    />
+  )
+}
+
 const MessagePart = {
   Text: MessagePartText,
   Dynamic: MessagePartDynamicTool,
@@ -270,6 +297,7 @@ const MessagePart = {
   QueryLogs: MessagePartQueryLogs,
   DeployEdgeFunction: MessagePartDeployEdgeFunction,
   NotebookProposal: MessagePartNotebookProposal,
+  NotebookRun: MessagePartNotebookRun,
 } as const
 
 function MessagePartContainer({
@@ -287,6 +315,7 @@ const isWideMessagePart = (part: NonNullable<VercelMessage['parts']>[number]) =>
   part.type === 'tool-query_logs' ||
   part.type === 'tool-create_notebook' ||
   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))
@@ -338,6 +367,9 @@ export function MessagePartSwitcher({
       case 'tool-update_notebook': {
         return <MessagePart.NotebookProposal toolPart={part} mode="update" />
       }
+      case 'tool-run_notebook': {
+        return <MessagePart.NotebookRun toolPart={part} />
+      }
 
       case 'source-url':
       case 'source-document':

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

@@ -141,6 +141,11 @@ export const updateNotebookInputSchema = z.object({
   operations: notebookOperationsSchema,
 })
 
+export const runNotebookInputSchema = z.object({
+  id: z.string(),
+  expected_updated_at: z.string(),
+})
+
 export const notebookToolOutputSchema = z.object({ id: z.string(), name: z.string() })
 
 export const updateNotebookToolOutputSchema = notebookToolOutputSchema.extend({

apps/studio/components/ui/AIAssistantPanel/NotebookRunRenderer.test.tsx

@@ -0,0 +1,211 @@
+import { screen } from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
+import { HttpResponse } from 'msw'
+import { describe, expect, it, vi } from 'vitest'
+
+import { NotebookRunRenderer } from './NotebookRunRenderer'
+import type { components } from '@/data/api'
+import { customRender as render } from '@/tests/lib/custom-render'
+import { addAPIMock, type APIErrorBody } from '@/tests/lib/msw'
+
+const NOTEBOOK_ID = 'd3aadd77-7c3c-4de7-aa5c-5aa8ac270b44'
+const UPDATED_AT = '2026-01-01T00:00:00.000Z'
+
+const mockNotebook = (updatedAt = UPDATED_AT) =>
+  addAPIMock({
+    method: 'get',
+    path: '/platform/projects/:ref/content/item/:id',
+    response: () =>
+      HttpResponse.json<components['schemas']['GetUserContentByIdResponse']>({
+        id: NOTEBOOK_ID,
+        type: 'notebook',
+        name: 'Signup funnel',
+        description: '',
+        favorite: false,
+        folder_id: null,
+        inserted_at: UPDATED_AT,
+        updated_at: updatedAt,
+        visibility: 'project',
+        owner_id: 1,
+        project_id: 1,
+        content: {
+          schema_version: 1,
+          cells: [
+            {
+              _tag: 'database_cell',
+              _id: 'cell-1',
+              title: 'Recent signups',
+              sql: 'select email from auth.users',
+              row_limit: 100,
+            },
+          ],
+        },
+      }),
+  })
+
+const mockNotebookError = () => {
+  addAPIMock({
+    method: 'get',
+    path: '/platform/projects/:ref',
+    response: {
+      id: 1,
+      ref: 'default',
+      organization_id: 1,
+      name: 'Test Project',
+      status: 'ACTIVE_HEALTHY',
+      cloud_provider: 'AWS',
+      region: 'us-east-1',
+      db_host: 'db.default.supabase.co',
+      restUrl: 'https://default.supabase.co/rest/v1/',
+      inserted_at: UPDATED_AT,
+      updated_at: UPDATED_AT,
+      subscription_id: 'sub-1',
+      is_branch_enabled: false,
+      is_physical_backups_enabled: false,
+      high_availability: false,
+      integration_source: null,
+      connectionString: 'postgresql://postgres@localhost:5432/postgres',
+      is_hibernating: false,
+    },
+  })
+  addAPIMock({
+    method: 'get',
+    path: '/platform/projects/:ref/content/item/:id',
+    response: () =>
+      HttpResponse.json<APIErrorBody>({ message: 'Notebook unavailable' }, { status: 500 }),
+  })
+}
+
+describe('NotebookRunRenderer', () => {
+  it('previews the notebook and requests one Run notebook approval', async () => {
+    const user = userEvent.setup()
+    const onApprove = vi.fn()
+    mockNotebook()
+
+    const { container } = render(
+      <NotebookRunRenderer
+        state="approval-requested"
+        confirmState="approval-requested"
+        input={{ id: NOTEBOOK_ID, expected_updated_at: UPDATED_AT }}
+        output={undefined}
+        onApprove={onApprove}
+        onDeny={vi.fn()}
+      />
+    )
+
+    const loadingStatus = container.querySelector('[aria-live="polite"]')
+    expect(loadingStatus).toHaveTextContent('Loading notebook...')
+    expect(loadingStatus?.querySelector('svg')).toHaveClass('motion-reduce:animate-none')
+    expect(await screen.findByText('Assistant wants to run "Signup funnel"')).toBeInTheDocument()
+    expect(container.querySelector('[aria-live="polite"]')).toBe(loadingStatus)
+    expect(loadingStatus).toHaveClass('sr-only')
+    expect(screen.getByText('1 cell')).toBeInTheDocument()
+    await user.click(screen.getByRole('button', { name: 'Run notebook' }))
+    expect(onApprove).toHaveBeenCalledTimes(1)
+  })
+
+  it('keeps raw results visible to the user after execution', async () => {
+    mockNotebook()
+
+    render(
+      <NotebookRunRenderer
+        state="output-available"
+        confirmState="success"
+        input={{ id: NOTEBOOK_ID, expected_updated_at: UPDATED_AT }}
+        output={{
+          id: NOTEBOOK_ID,
+          name: 'Signup funnel',
+          updated_at: UPDATED_AT,
+          cells: [
+            {
+              cell_id: 'cell-1',
+              title: 'Recent signups',
+              source: 'database',
+              status: 'success',
+              rows: [{ email: 'person@example.com' }],
+            },
+          ],
+        }}
+      />
+    )
+
+    expect(await screen.findByText('Notebook executed')).toBeInTheDocument()
+    expect(screen.getByText('1 row')).toBeInTheDocument()
+    expect(screen.getByText(/person@example\.com/)).toBeInTheDocument()
+  })
+
+  it('warns before approval when the notebook changed since the Assistant read it', async () => {
+    mockNotebook('2026-01-02T00:00:00.000Z')
+
+    render(
+      <NotebookRunRenderer
+        state="approval-requested"
+        confirmState="approval-requested"
+        input={{ id: NOTEBOOK_ID, expected_updated_at: UPDATED_AT }}
+        output={undefined}
+        onApprove={vi.fn()}
+        onDeny={vi.fn()}
+      />
+    )
+
+    expect(
+      await screen.findByText('Notebook changed since the Assistant read it')
+    ).toBeInTheDocument()
+    expect(screen.getByRole('button', { name: 'Run notebook' })).toBeInTheDocument()
+  })
+
+  it('keeps a failed approval request denyable without allowing the run', async () => {
+    const user = userEvent.setup()
+    const onApprove = vi.fn()
+    const onDeny = vi.fn()
+    mockNotebookError()
+
+    const { container } = render(
+      <NotebookRunRenderer
+        state="approval-requested"
+        confirmState="approval-requested"
+        input={{ id: NOTEBOOK_ID, expected_updated_at: UPDATED_AT }}
+        output={undefined}
+        onApprove={onApprove}
+        onDeny={onDeny}
+      />
+    )
+
+    expect(await screen.findByText('Failed to load notebook')).toBeInTheDocument()
+    expect(container.querySelector('[data-slot="assistant-confirm"]')).toBeInTheDocument()
+    expect(screen.queryByRole('button', { name: 'Run notebook' })).not.toBeInTheDocument()
+
+    await user.click(screen.getByRole('button', { name: 'Skip' }))
+    expect(onDeny).toHaveBeenCalledTimes(1)
+    expect(onApprove).not.toHaveBeenCalled()
+  })
+
+  it('warns when historical results are shown with a newer notebook', async () => {
+    mockNotebook('2026-01-02T00:00:00.000Z')
+
+    render(
+      <NotebookRunRenderer
+        state="output-available"
+        confirmState="success"
+        input={{ id: NOTEBOOK_ID, expected_updated_at: UPDATED_AT }}
+        output={{
+          id: NOTEBOOK_ID,
+          name: 'Signup funnel',
+          updated_at: UPDATED_AT,
+          cells: [
+            {
+              cell_id: 'cell-1',
+              title: 'Recent signups',
+              source: 'database',
+              status: 'success',
+              rows: [{ email: 'person@example.com' }],
+            },
+          ],
+        }}
+      />
+    )
+
+    expect(await screen.findByText('Notebook changed since this run')).toBeInTheDocument()
+    expect(screen.getByText(/preview shows the current notebook/i)).toBeInTheDocument()
+  })
+})

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

@@ -0,0 +1,202 @@
+import { useParams } from 'common'
+import { Loader2 } from 'lucide-react'
+import { Button, cn } from 'ui'
+import { Admonition } from 'ui-patterns/Admonition'
+import { CodeBlock } from 'ui-patterns/CodeBlock'
+
+import { AssistantNotebookPreview } from './AssistantNotebookPreview'
+import { toAssistantQueryResult } from './AssistantQueryCell.utils'
+import { Confirm } from './Confirm'
+import type { ConfirmFooterApprovalState } from './Confirm.utils'
+import { runNotebookInputSchema } from './Message.utils'
+import { AlertError } from '@/components/ui/AlertError'
+import { useNotebookQuery } from '@/data/content/notebooks/notebook-query'
+import { toWireNotebook } from '@/data/content/notebooks/notebook-schema'
+import { notebookRunOutputSchema } from '@/lib/ai/tools/notebook-run-output'
+
+export type NotebookRunState =
+  | 'input-available'
+  | 'approval-requested'
+  | 'approval-responded'
+  | 'output-denied'
+  | 'output-available'
+  | 'output-error'
+
+export interface NotebookRunRendererProps {
+  state: NotebookRunState
+  input: unknown
+  output: unknown
+  confirmState?: ConfirmFooterApprovalState
+  onApprove?: () => void
+  onDeny?: () => void
+}
+
+export const NotebookRunRenderer = ({
+  state,
+  input,
+  output,
+  confirmState,
+  onApprove,
+  onDeny,
+}: NotebookRunRendererProps) => {
+  const { ref } = useParams()
+  const parsedInput = runNotebookInputSchema.safeParse(input)
+  const {
+    data: notebook,
+    isLoading,
+    isError,
+    error,
+  } = useNotebookQuery(
+    { projectRef: ref, id: parsedInput.success ? parsedInput.data.id : undefined },
+    { enabled: parsedInput.success }
+  )
+
+  if (!parsedInput.success) {
+    return (
+      <Confirm
+        className="my-4"
+        state={confirmState}
+        message="Assistant wants to run a notebook"
+        denyOnly
+        onCancel={onDeny}
+      >
+        <div className="flex flex-col gap-2 p-3">
+          <Admonition
+            type="warning"
+            title="Couldn't render this notebook run"
+            description="The assistant's input didn't match the expected shape. You can review the raw input below."
+          />
+          <CodeBlock
+            language="json"
+            value={JSON.stringify(input, null, 2)}
+            hideLineNumbers
+            className="text-xs"
+            wrapperClassName="max-h-56"
+          />
+        </div>
+      </Confirm>
+    )
+  }
+
+  const loadingStatus = (
+    <div
+      aria-live="polite"
+      className={cn(
+        isLoading
+          ? 'my-4 rounded-lg border bg-surface-75 heading-meta h-9 px-3 text-foreground-light flex items-center gap-2'
+          : 'sr-only'
+      )}
+    >
+      {isLoading && (
+        <>
+          <Loader2 className="w-4 h-4 animate-spin motion-reduce:animate-none" />
+          Loading notebook...
+        </>
+      )}
+    </div>
+  )
+
+  if (isLoading) {
+    return <>{loadingStatus}</>
+  }
+
+  if (isError || !notebook) {
+    const loadError = <AlertError error={error} subject="Failed to load notebook" />
+
+    if (confirmState === 'approval-requested') {
+      return (
+        <>
+          {loadingStatus}
+          <Confirm
+            className="my-4"
+            state={confirmState}
+            message="Assistant wants to run a notebook"
+            denyOnly
+            onCancel={onDeny}
+          >
+            <div className="p-3">{loadError}</div>
+          </Confirm>
+        </>
+      )
+    }
+
+    return (
+      <>
+        {loadingStatus}
+        <div className="my-4 flex flex-col gap-2">
+          {loadError}
+          {confirmState !== undefined && (
+            <Button variant="outline" size="tiny" className="w-fit" disabled onClick={onDeny}>
+              Skip
+            </Button>
+          )}
+        </div>
+      </>
+    )
+  }
+
+  const entries = toWireNotebook(notebook.content).cells.map((cell) => ({
+    _tag: 'unchanged' as const,
+    cell,
+  }))
+  const parsedOutput = notebookRunOutputSchema.safeParse(output)
+  const isHistoricalRun = state === 'output-available' || state === 'output-error'
+  const referencedUpdatedAt =
+    state === 'output-available' && parsedOutput.success
+      ? parsedOutput.data.updated_at
+      : parsedInput.data.expected_updated_at
+  const hasNotebookChanged = notebook.updated_at !== referencedUpdatedAt
+  const results = parsedOutput.success
+    ? Object.fromEntries(
+        parsedOutput.data.cells.map((cell) => [
+          cell.cell_id,
+          cell.status === 'error'
+            ? { error: { message: cell.error?.message ?? 'Failed to run query' } }
+            : (toAssistantQueryResult(cell.rows ?? []) ?? { rows: [] }),
+        ])
+      )
+    : undefined
+
+  return (
+    <>
+      {loadingStatus}
+      <Confirm
+        className="my-4"
+        state={confirmState}
+        message={`Assistant wants to run "${notebook.name}"`}
+        cancelLabel="Skip"
+        confirmLabel="Run notebook"
+        confirmLabelLoading="Running..."
+        successMessage="Notebook executed"
+        errorMessage="Failed to run notebook"
+        deniedMessage="Skipped notebook run"
+        onCancel={onDeny}
+        onConfirm={onApprove}
+      >
+        {hasNotebookChanged && (
+          <div className="px-2 pt-2">
+            <Admonition
+              type="warning"
+              title={
+                isHistoricalRun
+                  ? 'Notebook changed since this run'
+                  : 'Notebook changed since the Assistant read it'
+              }
+              description={
+                isHistoricalRun
+                  ? 'This preview shows the current notebook. Its cells may not match the saved results from this run.'
+                  : 'Review the current cells below. The run will be rejected until the Assistant reads the latest notebook version.'
+              }
+            />
+          </div>
+        )}
+        <AssistantNotebookPreview
+          entries={entries}
+          mode="run"
+          title={notebook.name}
+          results={state === 'output-available' ? results : undefined}
+        />
+      </Confirm>
+    </>
+  )
+}