refactor(studio): migrate create, delete, and run notebook tools to tool contracts (#51754)

## Summary

Migrates the create, delete, and run notebook assistant tools to the
tool contract + server pattern used by `update_notebook`. Each tool now
defines its schema once in a shared contract, with server-only execution
split into a `.server.ts` module. This removes duplicated client-side
schemas and gives the three tools one shared notebook mutation output
schema.

## Changes

- Add `create-notebook`, `delete-notebook`, and `run-notebook` tool
contracts, each with a `.server.ts` implementation.
- Remove the duplicated client-side schemas for these tools from
`notebook-tools.ts` and the registry.
- Add `notebook-mutation-output.ts` as the single output schema shared
by the notebook mutation tools.
- Add `lib/ai/opt-in-level.ts` with the opt-in level helpers, extracted
from the existing module to break an import cycle.
- Add optional `ToolContract.outputOnInvalid` so `run_notebook` history
replay can reuse a stored output when input fails validation.
- Update the AI assistant panel renderers (proposal, run, opt-in modal
and request) and the tool filter, sanitizer, and cache invalidation to
use the new contracts.

## Deviations worth reviewing

- **`outputOnInvalid` on `ToolContract`.** This is a new optional hook.
It exists so `run_notebook` history replay can return a valid output
when stored input no longer passes validation. Reviewers should confirm
this is the right level for the hook rather than handling it inside the
tool.
- **`opt-in-level.ts` extraction.** The helpers were moved into their
own module to break a circular import between the tool modules and the
opt-in components. This adds one file, but no behavior change is
intended.
- **`create-notebook` contract imports the data-layer
`notebook-schema`.** This matches `update-notebook`. It should be moved
behind the contract before the planned ESLint rule step, which restricts
contract modules from importing data-layer code.

## Testing

- `vitest` for `lib/ai`, `components/ui/AIAssistantPanel`, and `state`:
920 tests pass across 71 files.
- `tsc`: one error remains in `packages/ui-patterns`
(`InstructionBlocks.tsx`), which is unrelated to this change. I have not
verified whether it also fails on the base branch.
- ESLint reports no errors on the changed files (9 `no-explicit-any`
warnings in `tool-filter.ts` and its test); Prettier is clean.
- Not run: knip and the lint ratchet.

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

@@ -19,7 +19,7 @@ import { useAIOptInForm } from '@/hooks/forms/useAIOptInForm'
 import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
 import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization'
-import { isOptInLevelAtLeast } from '@/lib/ai/tool-filter'
+import { isOptInLevelAtLeast } from '@/lib/ai/opt-in-level'
 
 interface AIOptInModalProps {
   visible: boolean

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

@@ -1,8 +1,6 @@
 import { untrustedSql } from '@supabase/pg-meta'
 import { z, type SafeParseReturnType } from 'zod'
 
-import { agentNotebookSchema } from '@/data/content/notebooks/notebook-schema'
-
 // Splits markdown into alternating [plain, code, plain, code, ...] segments.
 // Odd-indexed segments are already inside code spans/fences and should be left alone.
 const CODE_SEGMENT_REGEX = /(```[\s\S]*?```|`[^`]*`)/g
@@ -128,23 +126,6 @@ export const deployEdgeFunctionOutputSchema = z
   .object({ success: z.boolean().optional() })
   .passthrough()
 
-export const createNotebookInputSchema = z.object({
-  name: z.string(),
-  description: z.string().optional(),
-  content: agentNotebookSchema,
-})
-
-export const deleteNotebookInputSchema = z.object({
-  id: z.string(),
-})
-
-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 rateMessageResponseSchema = z.object({
   category: z.enum([
     'sql_generation',

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

@@ -9,11 +9,6 @@ import { CodeBlock } from 'ui-patterns/CodeBlock'
 import { AssistantNotebookPreview } from './AssistantNotebookPreview'
 import { Confirm } from './Confirm'
 import { type ConfirmFooterApprovalState } from './Confirm.utils'
-import {
-  createNotebookInputSchema,
-  deleteNotebookInputSchema,
-  notebookToolOutputSchema,
-} from './Message.utils'
 import { AlertError } from '@/components/ui/AlertError'
 import { hasDiscardableChanges } from '@/data/content/notebooks/notebook-cache'
 import {
@@ -23,6 +18,8 @@ import {
 } from '@/data/content/notebooks/notebook-operations'
 import { useNotebookQuery } from '@/data/content/notebooks/notebook-query'
 import { toWireNotebook } from '@/data/content/notebooks/notebook-schema'
+import { contract as createNotebookContract } from '@/lib/ai/tools/create-notebook'
+import { contract as deleteNotebookContract } from '@/lib/ai/tools/delete-notebook'
 import { contract as updateNotebookContract } from '@/lib/ai/tools/update-notebook'
 import { useNotebooksStateSnapshot } from '@/state/notebooks/notebooks-state'
 
@@ -96,6 +93,12 @@ const NOTEBOOK_ACTION_NOUN: Record<NotebookProposalMode, string> = {
  * fetch, and the `expected_updated_at` version check. Wraps the preview in `Confirm` the
  * same way `AssistantQueryCell` wraps `QueryEditor`.
  */
+const MODE_CONTRACTS = {
+  create: createNotebookContract,
+  update: updateNotebookContract,
+  delete: deleteNotebookContract,
+} as const satisfies Record<NotebookProposalMode, unknown>
+
 export const NotebookProposalRenderer = (props: NotebookProposalRendererProps) => {
   const { ref } = useParams()
   const { mode, state, input, output, confirmState, onApprove, onDeny, denyWithReason } = props
@@ -107,7 +110,7 @@ export const NotebookProposalRenderer = (props: NotebookProposalRendererProps) =
     return <UpdateNotebookTerminalSummary state={state} output={output} />
   }
 
-  const parsedOutput = notebookToolOutputSchema.safeParse(output)
+  const parsedOutput = MODE_CONTRACTS[mode].outputSchema.safeParse(output)
   // A deleted notebook no longer exists to open, so this action only applies to create/update.
   const canOpenNotebook = mode !== 'delete' && state === 'output-available'
   const footerAction =
@@ -163,7 +166,7 @@ function UpdateNotebookTerminalSummary({
   output,
 }: Pick<NotebookProposalRendererProps, 'state' | 'output'>) {
   const { ref } = useParams()
-  const parsedOutput = notebookToolOutputSchema.safeParse(output)
+  const parsedOutput = updateNotebookContract.outputSchema.safeParse(output)
   const label =
     state === 'output-available'
       ? parsedOutput.success
@@ -271,7 +274,7 @@ function CreateNotebookProposal({
   onApprove,
   onDeny,
 }: NotebookProposalStepProps) {
-  const parsedInput = createNotebookInputSchema.safeParse(input)
+  const parsedInput = createNotebookContract.inputSchema.safeParse(input)
 
   if (!parsedInput.success) {
     return (
@@ -523,7 +526,7 @@ function DeleteNotebookProposal({
 }: NotebookProposalStepPropsWithOutput) {
   const { ref } = useParams()
   const notebooksSnap = useNotebooksStateSnapshot()
-  const parsedInput = deleteNotebookInputSchema.safeParse(input)
+  const parsedInput = deleteNotebookContract.inputSchema.safeParse(input)
   const isCompleted = state === 'output-available'
   const hasUnsavedLocalChanges =
     confirmState === 'approval-requested' &&
@@ -552,7 +555,7 @@ function DeleteNotebookProposal({
   }
 
   if (isCompleted) {
-    const parsedOutput = notebookToolOutputSchema.safeParse(output)
+    const parsedOutput = deleteNotebookContract.outputSchema.safeParse(output)
     const notebookName = parsedOutput.success ? parsedOutput.data.name : undefined
 
     return (

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

@@ -8,11 +8,10 @@ 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'
+import { contract as runNotebookContract } from '@/lib/ai/tools/run-notebook'
 
 export type NotebookRunState =
   | 'input-available'
@@ -40,7 +39,7 @@ export const NotebookRunRenderer = ({
   onDeny,
 }: NotebookRunRendererProps) => {
   const { ref } = useParams()
-  const parsedInput = runNotebookInputSchema.safeParse(input)
+  const parsedInput = runNotebookContract.inputSchema.safeParse(input)
   const {
     data: notebook,
     isLoading,
@@ -139,7 +138,7 @@ export const NotebookRunRenderer = ({
     _tag: 'unchanged' as const,
     cell,
   }))
-  const parsedOutput = notebookRunOutputSchema.safeParse(output)
+  const parsedOutput = runNotebookContract.outputSchema.safeParse(output)
   const isHistoricalRun = state === 'output-available' || state === 'output-error'
   const referencedUpdatedAt =
     state === 'output-available' && parsedOutput.success

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

@@ -9,7 +9,7 @@ import type { ConfirmFooterApprovalState } from './Confirm.utils'
 import { NoPermission } from '@/components/ui/NoPermission'
 import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions'
 import { useOrgAiOptInLevel, type AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
-import { optInLevelSchema } from '@/lib/ai/tool-filter'
+import { optInLevelSchema } from '@/lib/ai/opt-in-level'
 
 const updateOptInLevelOutputSchema = z.object({
   levelAfterReview: optInLevelSchema,

apps/studio/eslint.config.cjs

@@ -69,6 +69,14 @@ const TOOL_MODULE_RESTRICTED_IMPORT_PATHS = [
     message: APPROVAL_GATE_MESSAGE,
   },
 ]
+// The gate itself is also off limits: only `approvalGatedToolFromContract` may hand it out.
+const TOOL_MODULE_RESTRICTED_IMPORT_PATTERNS = [
+  {
+    group: ['**/approval-gate', './approval-gate'],
+    importNames: ['approvalGate'],
+    message: APPROVAL_GATE_MESSAGE,
+  },
+]
 
 module.exports = defineConfig([
   { files: ['**/*.ts', '**/*.tsx'] },
@@ -145,7 +153,7 @@ module.exports = defineConfig([
         'error',
         {
           paths: [...BASE_RESTRICTED_IMPORT_PATHS, ...TOOL_MODULE_RESTRICTED_IMPORT_PATHS],
-          patterns: LUCIDE_ICON_IMPORT_PATTERNS,
+          patterns: [...LUCIDE_ICON_IMPORT_PATTERNS, ...TOOL_MODULE_RESTRICTED_IMPORT_PATTERNS],
         },
       ],
     },

apps/studio/evals/scorer-wasm.ts

@@ -3,7 +3,7 @@ import { parse } from 'libpg-query'
 
 import { AssistantEvalScorer } from './scorer'
 import { getParsedToolSpans } from './trace-utils'
-import { createNotebookInputSchema } from '@/components/ui/AIAssistantPanel/Message.utils'
+import { contract as createNotebookContract } from '@/lib/ai/tools/create-notebook'
 import { executeSqlInputSchema } from '@/lib/ai/tools/studio-tools'
 import { extractIdentifiers, isQuotedInSql, needsQuoting } from '@/lib/sql-identifier-quoting'
 
@@ -17,7 +17,7 @@ async function getSqlQueries(trace: Trace): Promise<string[]> {
     inputSchema: executeSqlInputSchema,
   })
   const createNotebookSpans = await getParsedToolSpans(trace, 'create_notebook', {
-    inputSchema: createNotebookInputSchema,
+    inputSchema: createNotebookContract.inputSchema,
   })
 
   const notebookCellSql = createNotebookSpans.flatMap((s) =>

apps/studio/lib/ai/notebook-cache-invalidation.ts

@@ -1,10 +1,9 @@
 import type { QueryClient } from '@tanstack/react-query'
 import type { ToolUIPart, UIMessage } from 'ai'
 
-import { notebookToolOutputSchema } from '@/components/ui/AIAssistantPanel/Message.utils'
 import { contentKeys } from '@/data/content/keys'
 import { evictNotebookFromCaches } from '@/data/content/notebooks/notebook-cache'
-import type { ToolName } from '@/lib/ai/tool-filter'
+import { TOOL_CONTRACTS, type ContractToolName } from '@/lib/ai/tools/registry'
 import { notebooksState } from '@/state/notebooks/notebooks-state'
 
 export type NotebookCacheEffect =
@@ -15,13 +14,21 @@ const NOTEBOOK_MUTATION_TOOLS = [
   'create_notebook',
   'update_notebook',
   'delete_notebook',
-] as const satisfies readonly ToolName[]
+] as const satisfies readonly ContractToolName[]
+
+const NOTEBOOK_MUTATION_CONTRACTS = {
+  'tool-create_notebook': TOOL_CONTRACTS.create_notebook,
+  'tool-update_notebook': TOOL_CONTRACTS.update_notebook,
+  'tool-delete_notebook': TOOL_CONTRACTS.delete_notebook,
+} as const
 
 const NOTEBOOK_MUTATION_TOOL_TYPES: ReadonlySet<string> = new Set(
   NOTEBOOK_MUTATION_TOOLS.map((name) => `tool-${name}`)
 )
 
-function isNotebookMutationPart(part: UIMessage['parts'][number]): part is ToolUIPart {
+function isNotebookMutationPart(
+  part: UIMessage['parts'][number]
+): part is ToolUIPart & { type: keyof typeof NOTEBOOK_MUTATION_CONTRACTS } {
   return NOTEBOOK_MUTATION_TOOL_TYPES.has(part.type)
 }
 
@@ -39,7 +46,7 @@ export function collectNotebookCacheEffects(
       if (part.state !== 'output-available') continue
       if (processed.has(part.toolCallId)) continue
 
-      const result = notebookToolOutputSchema.safeParse(part.output)
+      const result = NOTEBOOK_MUTATION_CONTRACTS[part.type].outputSchema.safeParse(part.output)
       if (!result.success) continue
 
       effects.push({

apps/studio/lib/ai/opt-in-level.ts

@@ -0,0 +1,21 @@
+import { z } from 'zod'
+
+import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
+
+/** Ordered from least to most data shared with the model. */
+export const OPT_IN_LEVELS = [
+  'disabled',
+  'schema',
+  'schema_and_log',
+  'schema_and_log_and_data',
+] as const
+
+export const optInLevelSchema = z.enum(OPT_IN_LEVELS)
+
+export function isOptInLevelAtLeast(level: AiOptInLevel, minimum: AiOptInLevel): boolean {
+  return OPT_IN_LEVELS.indexOf(level) >= OPT_IN_LEVELS.indexOf(minimum)
+}
+
+export function lowerOptInLevel(a: AiOptInLevel, b: AiOptInLevel): AiOptInLevel {
+  return isOptInLevelAtLeast(a, b) ? b : a
+}

apps/studio/lib/ai/tool-filter.test.ts

@@ -9,31 +9,44 @@ import {
   TOOL_CATEGORY_MAP,
   toolSetValidationSchema,
 } from './tool-filter'
+import { contract as createNotebookContract } from './tools/create-notebook'
+import { contract as deleteNotebookContract } from './tools/delete-notebook'
+import { contract as runNotebookContract } from './tools/run-notebook'
 import { contract as updateNotebookContract } from './tools/update-notebook'
 
 describe('TOOL_CATEGORY_MAP', () => {
   it('should categorize tools correctly', () => {
     expect(TOOL_CATEGORY_MAP['execute_sql']).toBe(TOOL_CATEGORIES.UI)
     expect(TOOL_CATEGORY_MAP['list_tables']).toBe(TOOL_CATEGORIES.SCHEMA)
-    expect(TOOL_CATEGORY_MAP['run_notebook']).toBe(TOOL_CATEGORIES.SCHEMA)
+    expect(TOOL_CATEGORY_MAP['list_notebooks']).toBe(TOOL_CATEGORIES.SCHEMA)
   })
 })
 
 describe('tools with a contract', () => {
-  const tools = {
-    update_notebook: { description: 'Update a notebook', execute: vitest.fn() },
-  } as unknown as ToolSet
+  const contracts = {
+    create_notebook: createNotebookContract,
+    delete_notebook: deleteNotebookContract,
+    run_notebook: runNotebookContract,
+    update_notebook: updateNotebookContract,
+  }
 
-  it('is gated by the contract minOptInLevel, not the category map', () => {
-    expect(updateNotebookContract.minOptInLevel).toBe('schema')
-    expect(Object.hasOwn(TOOL_CATEGORY_MAP, 'update_notebook')).toBe(false)
+  it.each(Object.entries(contracts))(
+    '%s is gated by the contract minOptInLevel, not the category map',
+    (name, contract) => {
+      const tools = {
+        [name]: { description: 'A notebook tool', execute: vitest.fn() },
+      } as unknown as ToolSet
 
-    const blocked = filterToolsByOptInLevel(tools, 'disabled').update_notebook
-    const allowed = filterToolsByOptInLevel(tools, 'schema').update_notebook
+      expect(contract.minOptInLevel).toBe('schema')
+      expect(Object.hasOwn(TOOL_CATEGORY_MAP, name)).toBe(false)
 
-    expect(blocked.description).toContain('Requires opting in')
-    expect(allowed.description).toBe('Update a notebook')
-  })
+      const blocked = filterToolsByOptInLevel(tools, 'disabled')[name]
+      const allowed = filterToolsByOptInLevel(tools, 'schema')[name]
+
+      expect(blocked.description).toContain('Requires opting in')
+      expect(allowed.description).toBe('A notebook tool')
+    }
+  )
 })
 
 describe('tool allowance by opt-in level', () => {

apps/studio/lib/ai/tool-filter.ts

@@ -1,6 +1,7 @@
 import type { Tool, ToolSet } from 'ai'
 import { z } from 'zod'
 
+import { isOptInLevelAtLeast, optInLevelSchema } from './opt-in-level'
 import { getToolContract, type ContractToolName } from './tools/registry'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
 
@@ -60,16 +61,6 @@ export const toolSetValidationSchema = z.record(
 )
 export type ToolName = keyof z.infer<typeof toolSetValidationSchema>
 
-/** Ordered from least to most data shared with the model. */
-export const OPT_IN_LEVELS = [
-  'disabled',
-  'schema',
-  'schema_and_log',
-  'schema_and_log_and_data',
-] as const
-
-export const optInLevelSchema = z.enum(OPT_IN_LEVELS)
-
 export const updateOptInLevelInputSchema = z.object({
   requiredLevel: z
     .enum(['schema', 'schema_and_log', 'schema_and_log_and_data'])
@@ -84,14 +75,6 @@ export const storedUpdateOptInLevelInputSchema = updateOptInLevelInputSchema.ext
   levelWhenAsked: optInLevelSchema.optional(),
 })
 
-export function isOptInLevelAtLeast(level: AiOptInLevel, minimum: AiOptInLevel): boolean {
-  return OPT_IN_LEVELS.indexOf(level) >= OPT_IN_LEVELS.indexOf(minimum)
-}
-
-export function lowerOptInLevel(a: AiOptInLevel, b: AiOptInLevel): AiOptInLevel {
-  return isOptInLevelAtLeast(a, b) ? b : a
-}
-
 /**
  * Tool categories based on the data they access
  */
@@ -135,9 +118,6 @@ export const TOOL_CATEGORY_MAP: Record<Exclude<ToolName, ContractToolName>, Tool
   list_databases: TOOL_CATEGORIES.SCHEMA,
   list_notebooks: TOOL_CATEGORIES.SCHEMA,
   get_notebook: TOOL_CATEGORIES.SCHEMA,
-  run_notebook: TOOL_CATEGORIES.SCHEMA,
-  create_notebook: TOOL_CATEGORIES.SCHEMA,
-  delete_notebook: TOOL_CATEGORIES.SCHEMA,
   getSchemaTables: TOOL_CATEGORIES.SCHEMA,
   getRlsKnowledge: TOOL_CATEGORIES.SCHEMA,
   getFunctions: TOOL_CATEGORIES.SCHEMA,

apps/studio/lib/ai/tools/create-notebook.server.ts

@@ -0,0 +1,61 @@
+import { untrustedSql } from '@supabase/pg-meta'
+
+import { getAuthHeaders } from './auth-headers'
+import { contract } from './create-notebook'
+import { assertValidDatabaseIdentifiers } from './notebook-tool-error'
+import { approvalGatedToolFromContract } from './tool-from-contract'
+import type { ApprovalGatedToolFactory } from './types'
+import type { WritableCell, WritableNotebook } from '@/data/content/notebooks/notebook-schema'
+import { createNotebook } from '@/data/content/notebooks/notebook-upsert-mutation'
+import { untrustedLogSql } from '@/data/logs/safe-analytics-sql'
+import { getReadReplicas } from '@/data/read-replicas/replicas-query'
+
+export const createTool: ApprovalGatedToolFactory<typeof contract> = (ctx) => {
+  const { projectRef } = ctx
+  const authHeaders = getAuthHeaders(ctx.authorization)
+
+  return approvalGatedToolFromContract(
+    contract,
+    ctx,
+    async ({ name, description, content }, { abortSignal }, gate) => {
+      if (
+        content.cells.some(
+          (cell) => cell._tag === 'database_cell' && cell.database_identifier !== undefined
+        )
+      ) {
+        const databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders)
+        assertValidDatabaseIdentifiers(
+          content.cells,
+          new Set((databases ?? []).map((database) => database.identifier))
+        )
+      }
+
+      // `gate` is only handed to an `execute` that runs after the user approves, so that
+      // approval is the user gesture that promotes each cell's SQL from untrusted to safe.
+      // Keep the promotion here, not in a shared helper, so it's auditable directly.
+      const cells: WritableNotebook['cells'] = content.cells.map((cell): WritableCell => {
+        switch (cell._tag) {
+          case 'markdown_cell':
+            return cell
+          case 'database_cell':
+            return { ...cell, sql: gate.acceptUntrustedSql(untrustedSql(cell.sql)) }
+          case 'log_cell':
+            return { ...cell, sql: gate.acceptUntrustedLogsSql(untrustedLogSql(cell.sql)) }
+        }
+      })
+
+      const result = await createNotebook(
+        {
+          projectRef: projectRef ?? '',
+          name,
+          description,
+          content: { schema_version: content.schema_version, cells },
+        },
+        abortSignal,
+        authHeaders
+      )
+
+      return { id: result.id, name }
+    }
+  )
+}

apps/studio/lib/ai/tools/create-notebook.ts

@@ -0,0 +1,24 @@
+import { z } from 'zod'
+
+import { notebookMutationOutputSchema } from './notebook-mutation-output'
+import type { ToolContract } from './types'
+import { agentNotebookSchema } from '@/data/content/notebooks/notebook-schema'
+
+const inputSchema = z.object({
+  name: z.string().describe('A short, descriptive name for the notebook.'),
+  description: z.string().optional().describe('A short description of what the notebook is for.'),
+  content: agentNotebookSchema.describe(
+    'The notebook content: a schema version and an ordered list of cells (markdown, database, or log). Cells must not include an id — one is assigned when the notebook is saved.'
+  ),
+})
+
+export const contract = {
+  name: 'create_notebook',
+  description:
+    'Asks the user to create a new notebook with the given cells. Requires user approval before creating.',
+  inputSchema,
+  outputSchema: notebookMutationOutputSchema,
+  minOptInLevel: 'schema',
+  needsApproval: true,
+  toModelOutput: ({ id, name }) => ({ id, name }),
+} as const satisfies ToolContract

apps/studio/lib/ai/tools/delete-notebook.server.ts

@@ -0,0 +1,19 @@
+import { getAuthHeaders } from './auth-headers'
+import { contract } from './delete-notebook'
+import { toolFromContract } from './tool-from-contract'
+import type { ToolFactory } from './types'
+import { deleteContents } from '@/data/content/content-delete-mutation'
+import { getNotebook } from '@/data/content/notebooks/notebook-query'
+
+export const createTool: ToolFactory<typeof contract> = (ctx) => {
+  const { projectRef } = ctx
+  const authHeaders = getAuthHeaders(ctx.authorization)
+
+  return toolFromContract(contract, ctx, async ({ id }, { abortSignal }) => {
+    const notebook = await getNotebook({ projectRef, id }, abortSignal, authHeaders)
+
+    await deleteContents({ projectRef: projectRef ?? '', ids: [id] }, abortSignal, authHeaders)
+
+    return { id, name: notebook.name }
+  })
+}

apps/studio/lib/ai/tools/delete-notebook.ts

@@ -0,0 +1,18 @@
+import { z } from 'zod'
+
+import { notebookMutationOutputSchema } from './notebook-mutation-output'
+import type { ToolContract } from './types'
+
+const inputSchema = z.object({
+  id: z.string().describe('The id of the notebook to delete.'),
+})
+
+export const contract = {
+  name: 'delete_notebook',
+  description: 'Asks the user to delete a notebook. Requires user approval before deleting.',
+  inputSchema,
+  outputSchema: notebookMutationOutputSchema,
+  minOptInLevel: 'schema',
+  needsApproval: true,
+  toModelOutput: ({ id, name }) => ({ id, name }),
+} as const satisfies ToolContract

apps/studio/lib/ai/tools/notebook-mutation-output.ts

@@ -0,0 +1,4 @@
+import { z } from 'zod'
+
+/** What a notebook mutation (create, update, delete) reports back: the notebook it touched. */
+export const notebookMutationOutputSchema = z.object({ id: z.string(), name: z.string() })

apps/studio/lib/ai/tools/notebook-run-output.ts

@@ -1,7 +1,7 @@
 import type { JSONValue } from 'ai'
 import { z } from 'zod'
 
-import { lowerOptInLevel, optInLevelSchema } from '../tool-filter'
+import { lowerOptInLevel, optInLevelSchema } from '../opt-in-level'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
 
 const jsonValueSchema: z.ZodType<JSONValue> = z.lazy(() =>

apps/studio/lib/ai/tools/notebook-tools.ts

@@ -1,46 +1,20 @@
-import { acceptUntrustedSql, untrustedSql } from '@supabase/pg-meta'
-import { tool, type JSONValue } from 'ai'
+import { tool } from 'ai'
 import { z } from 'zod'
 
 import { getAuthHeaders } from './auth-headers'
-import {
-  sanitizeNotebookRunOutput,
-  type NotebookRunCellOutput,
-  type NotebookRunOutput,
-} from './notebook-run-output'
-import { assertValidDatabaseIdentifiers, NotebookToolError } from './notebook-tool-error'
+import { createTool as createCreateNotebookTool } from './create-notebook.server'
+import { createTool as createDeleteNotebookTool } from './delete-notebook.server'
+import { createTool as createRunNotebookTool } from './run-notebook.server'
 import type { ToolContext } from './types'
 import { createTool as createUpdateNotebookTool } from './update-notebook.server'
-import { resolveLogTimeRange } from '@/components/interfaces/QuerySources/LogTimeRange.utils'
-import { deleteContents } from '@/data/content/content-delete-mutation'
 import { getContent } from '@/data/content/content-infinite-query'
 import { getNotebook } from '@/data/content/notebooks/notebook-query'
-import {
-  agentNotebookSchema,
-  toWireNotebook,
-  type WritableCell,
-  type WritableNotebook,
-} from '@/data/content/notebooks/notebook-schema'
-import { createNotebook } from '@/data/content/notebooks/notebook-upsert-mutation'
-import { executeLogsSql } from '@/data/logs/execute-logs-sql-mutation'
-import { acceptUntrustedLogsSql, untrustedLogSql } from '@/data/logs/safe-analytics-sql'
-import { QUERY_SOURCE_REGISTRY } from '@/data/query-sources/query-source-registry'
+import { toWireNotebook } from '@/data/content/notebooks/notebook-schema'
 import { getReadReplicas } from '@/data/read-replicas/replicas-query'
-import { executeSql } from '@/data/sql/execute-sql-mutation'
-import { applyAutoLimit } from '@/data/sql/utils'
 import type { Notebooks } from '@/types'
 
-function getErrorMessage(error: unknown): string {
-  if (error instanceof Error) return error.message
-  if (typeof error === 'object' && error !== null && 'message' in error) {
-    const message = (error as { message?: unknown }).message
-    if (typeof message === 'string') return message
-  }
-  return 'An unexpected error occurred while running this query.'
-}
-
 export const getNotebookTools = (ctx: ToolContext = {}) => {
-  const { projectRef, connectionString, authorization, aiOptInLevel = 'schema' } = ctx
+  const { projectRef, authorization } = ctx
   const authHeaders = getAuthHeaders(authorization)
 
   return {
@@ -123,211 +97,9 @@ export const getNotebookTools = (ctx: ToolContext = {}) => {
         }
       },
     }),
-    run_notebook: tool({
-      description:
-        'Asks the user to run every database and logs query cell in a notebook, in notebook order, and returns each result according to the organization data-sharing permissions. Requires one user approval before any cell executes.',
-      inputSchema: z.object({
-        id: z.string().describe('The id of the notebook to run.'),
-        expected_updated_at: z
-          .string()
-          .describe(
-            'The `updated_at` received from `get_notebook`. The run is rejected if the notebook changed since.'
-          ),
-      }),
-      needsApproval: true,
-      execute: async ({ id, expected_updated_at }, { abortSignal }) => {
-        const notebook = await getNotebook({ projectRef, id }, abortSignal, authHeaders)
-
-        if (notebook.updated_at !== expected_updated_at) {
-          throw new NotebookToolError(
-            `Notebook "${id}" changed since expected_updated_at (${expected_updated_at}); it is now ${notebook.updated_at}. Call get_notebook again and reissue run_notebook against the current content.`,
-            { exposeToAssistant: true }
-          )
-        }
-
-        const queryCells = notebook.content.cells.filter(
-          (cell) => cell._tag === 'database_cell' || cell._tag === 'log_cell'
-        )
-        const shouldLookupReplica = queryCells.some(
-          (cell) =>
-            cell._tag === 'database_cell' &&
-            cell.database_identifier !== undefined &&
-            cell.database_identifier !== projectRef
-        )
-        let databases: Awaited<ReturnType<typeof getReadReplicas>> = []
-        let replicaLookupFailure: { error: unknown } | undefined
-        if (shouldLookupReplica) {
-          try {
-            databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders)
-          } catch (error) {
-            replicaLookupFailure = { error }
-          }
-        }
-
-        const cells: NotebookRunCellOutput[] = []
-
-        // This execute function only runs after the AI SDK approval gate above resolves true.
-        // That approval is the explicit user gesture which promotes every persisted cell's
-        // untrusted SQL immediately before execution. Run sequentially to preserve notebook
-        // order and avoid racing cells that may depend on earlier writes.
-        for (const cell of queryCells) {
-          const title = cell.title?.trim() || 'Untitled query'
-
-          try {
-            if (cell._tag === 'log_cell') {
-              const result = await executeLogsSql({
-                projectRef: projectRef ?? '',
-                sql: acceptUntrustedLogsSql(cell.unchecked_sql),
-                range: resolveLogTimeRange(cell.time_range),
-                endpoint: QUERY_SOURCE_REGISTRY.logs.endpoint,
-                signal: abortSignal,
-                headers: authHeaders,
-              })
-              if (result.error) throw result.error
-              cells.push({
-                cell_id: cell._id,
-                title,
-                source: 'logs',
-                status: 'success',
-                // The analytics response crossed a JSON fetch boundary, although its shared
-                // data-layer type remains `unknown[]` because individual endpoints vary.
-                rows: result.rows as JSONValue[],
-              })
-              continue
-            }
-
-            const isReplicaTarget =
-              cell.database_identifier !== undefined && cell.database_identifier !== projectRef
-
-            if (isReplicaTarget && replicaLookupFailure !== undefined) {
-              throw replicaLookupFailure.error
-            }
-
-            const selectedConnectionString = !isReplicaTarget
-              ? connectionString
-              : databases?.find((database) => database.identifier === cell.database_identifier)
-                  ?.connectionString
-
-            if (!selectedConnectionString) {
-              throw new Error(
-                `Unable to run query: connection string is missing for database "${cell.database_identifier ?? projectRef ?? 'primary'}"`
-              )
-            }
-
-            const limitedSql = applyAutoLimit(
-              acceptUntrustedSql(cell.unchecked_sql),
-              cell.row_limit
-            )
-            const { result } = await executeSql<JSONValue[]>(
-              {
-                projectRef,
-                connectionString: selectedConnectionString,
-                sql: limitedSql.sql,
-                isStatementTimeoutDisabled: true,
-              },
-              abortSignal,
-              authHeaders
-            )
-            cells.push({
-              cell_id: cell._id,
-              title,
-              source: 'database',
-              status: 'success',
-              rows: Array.isArray(result) ? result : [],
-            })
-          } catch (error) {
-            cells.push({
-              cell_id: cell._id,
-              title,
-              source: cell._tag === 'log_cell' ? 'logs' : 'database',
-              status: 'error',
-              error: { message: getErrorMessage(error) },
-            })
-          }
-        }
-
-        return {
-          id: notebook.id,
-          name: notebook.name,
-          updated_at: notebook.updated_at,
-          optInLevel: aiOptInLevel,
-          cells,
-        } satisfies NotebookRunOutput
-      },
-      toModelOutput: ({ output }) => ({
-        type: 'json',
-        value: sanitizeNotebookRunOutput(output, aiOptInLevel),
-      }),
-    }),
-    create_notebook: tool({
-      description:
-        'Asks the user to create a new notebook with the given cells. Requires user approval before creating.',
-      inputSchema: z.object({
-        name: z.string().describe('A short, descriptive name for the notebook.'),
-        description: z
-          .string()
-          .optional()
-          .describe('A short description of what the notebook is for.'),
-        content: agentNotebookSchema.describe(
-          'The notebook content: a schema version and an ordered list of cells (markdown, database, or log). Cells must not include an id — one is assigned when the notebook is saved.'
-        ),
-      }),
-      needsApproval: true,
-      execute: async ({ name, description, content }, { abortSignal }) => {
-        if (
-          content.cells.some(
-            (cell) => cell._tag === 'database_cell' && cell.database_identifier !== undefined
-          )
-        ) {
-          const databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders)
-          assertValidDatabaseIdentifiers(
-            content.cells,
-            new Set((databases ?? []).map((database) => database.identifier))
-          )
-        }
-
-        // This approval gate is the user gesture that promotes each cell's SQL from
-        // untrusted to safe — keep the promotion here, not in a shared helper, so it's
-        // auditable directly alongside the `needsApproval: true` above.
-        const cells: WritableNotebook['cells'] = content.cells.map((cell): WritableCell => {
-          switch (cell._tag) {
-            case 'markdown_cell':
-              return cell
-            case 'database_cell':
-              return { ...cell, sql: acceptUntrustedSql(untrustedSql(cell.sql)) }
-            case 'log_cell':
-              return { ...cell, sql: acceptUntrustedLogsSql(untrustedLogSql(cell.sql)) }
-          }
-        })
-
-        const result = await createNotebook(
-          {
-            projectRef: projectRef ?? '',
-            name,
-            description,
-            content: { schema_version: content.schema_version, cells },
-          },
-          abortSignal,
-          authHeaders
-        )
-
-        return { id: result.id, name }
-      },
-    }),
+    run_notebook: createRunNotebookTool(ctx),
+    create_notebook: createCreateNotebookTool(ctx),
     update_notebook: createUpdateNotebookTool(ctx),
-    delete_notebook: tool({
-      description: 'Asks the user to delete a notebook. Requires user approval before deleting.',
-      inputSchema: z.object({
-        id: z.string().describe('The id of the notebook to delete.'),
-      }),
-      needsApproval: true,
-      execute: async ({ id }, { abortSignal }) => {
-        const notebook = await getNotebook({ projectRef, id }, abortSignal, authHeaders)
-
-        await deleteContents({ projectRef: projectRef ?? '', ids: [id] }, abortSignal, authHeaders)
-
-        return { id, name: notebook.name }
-      },
-    }),
+    delete_notebook: createDeleteNotebookTool(ctx),
   }
 }

apps/studio/lib/ai/tools/registry.test.ts

@@ -1,6 +1,7 @@
 import type { ToolUIPart } from 'ai'
 import { describe, expect, expectTypeOf, it } from 'vitest'
 
+import { INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE } from './notebook-run-output'
 import { sanitizeContractPart, type StudioUIMessage } from './registry'
 
 describe('TOOL_CONTRACTS', () => {
@@ -46,3 +47,59 @@ describe('sanitizeContractPart', () => {
     ).toBe(undefined)
   })
 })
+
+describe('sanitizeContractPart for run_notebook', () => {
+  const output = {
+    id: 'notebook-1',
+    name: 'Notebook',
+    updated_at: '2026-01-01T00:00:00.000Z',
+    optInLevel: 'schema_and_log_and_data',
+    cells: [
+      {
+        cell_id: 'db',
+        title: 'Database',
+        source: 'database',
+        status: 'success',
+        rows: [{ secret: 'value' }],
+      },
+    ],
+  }
+  const part = {
+    type: 'tool-run_notebook',
+    toolCallId: 'call-2',
+    state: 'output-available',
+    input: {},
+    output,
+  } as ToolUIPart
+
+  it('hides rows below the data level and keeps them at the data level', () => {
+    const hidden = sanitizeContractPart(part, 'schema')?.output as typeof output
+    expect(hidden.cells[0]).not.toHaveProperty('rows')
+    expect(sanitizeContractPart(part, 'schema_and_log_and_data')?.output).toEqual(output)
+  })
+
+  it('replaces an invalid output with an explanation instead of dropping it', () => {
+    expect(sanitizeContractPart({ ...part, output: 'oops' } as ToolUIPart, 'schema')?.output).toBe(
+      INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE
+    )
+  })
+})
+
+describe('notebook mutation contracts', () => {
+  it.each(['create_notebook', 'delete_notebook'] as const)(
+    '%s output is limited to id and name on replay',
+    (name) => {
+      const part = {
+        type: `tool-${name}`,
+        toolCallId: 'call-3',
+        state: 'output-available',
+        input: {},
+        output: { id: 'notebook-1', name: 'Notebook', extra: 'ignored' },
+      } as ToolUIPart
+      expect(sanitizeContractPart(part, 'schema')?.output).toEqual({
+        id: 'notebook-1',
+        name: 'Notebook',
+      })
+    }
+  )
+})

apps/studio/lib/ai/tools/registry.ts

@@ -1,6 +1,9 @@
 import type { InferUITools, Tool, ToolUIPart, UIDataTypes, UIMessage } from 'ai'
 
 import type { ToolName } from '../tool-filter'
+import { contract as createNotebook } from './create-notebook'
+import { contract as deleteNotebook } from './delete-notebook'
+import { contract as runNotebook } from './run-notebook'
 import type { ToolContract, ToolContractInput, ToolContractOutput } from './types'
 import { contract as updateNotebook } from './update-notebook'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
@@ -10,6 +13,9 @@ import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
  * by hand in the places the contract replaces.
  */
 export const TOOL_CONTRACTS = {
+  create_notebook: createNotebook,
+  delete_notebook: deleteNotebook,
+  run_notebook: runNotebook,
   update_notebook: updateNotebook,
 } as const satisfies Partial<Record<ToolName, ToolContract>>
 
@@ -53,6 +59,8 @@ export function sanitizeContractPart<Part extends ToolUIPart>(
   const parsed = contract.outputSchema.safeParse(part.output)
   return {
     ...part,
-    output: parsed.success ? contract.toModelOutput(parsed.data, optInLevel) : undefined,
+    output: parsed.success
+      ? contract.toModelOutput(parsed.data, optInLevel)
+      : contract.outputOnInvalid,
   }
 }

apps/studio/lib/ai/tools/run-notebook.server.ts

@@ -0,0 +1,153 @@
+import type { JSONValue } from 'ai'
+
+import { getAuthHeaders } from './auth-headers'
+import type { NotebookRunCellOutput, NotebookRunOutput } from './notebook-run-output'
+import { NotebookToolError } from './notebook-tool-error'
+import { contract } from './run-notebook'
+import { approvalGatedToolFromContract } from './tool-from-contract'
+import type { ApprovalGatedToolFactory } from './types'
+import { resolveLogTimeRange } from '@/components/interfaces/QuerySources/LogTimeRange.utils'
+import { getNotebook } from '@/data/content/notebooks/notebook-query'
+import { executeLogsSql } from '@/data/logs/execute-logs-sql-mutation'
+import { QUERY_SOURCE_REGISTRY } from '@/data/query-sources/query-source-registry'
+import { getReadReplicas } from '@/data/read-replicas/replicas-query'
+import { executeSql } from '@/data/sql/execute-sql-mutation'
+import { applyAutoLimit } from '@/data/sql/utils'
+
+function getErrorMessage(error: unknown): string {
+  if (error instanceof Error) return error.message
+  if (typeof error === 'object' && error !== null && 'message' in error) {
+    const message = (error as { message?: unknown }).message
+    if (typeof message === 'string') return message
+  }
+  return 'An unexpected error occurred while running this query.'
+}
+
+export const createTool: ApprovalGatedToolFactory<typeof contract> = (ctx) => {
+  const { projectRef, connectionString, aiOptInLevel = 'schema' } = ctx
+  const authHeaders = getAuthHeaders(ctx.authorization)
+
+  return approvalGatedToolFromContract(
+    contract,
+    ctx,
+    async ({ id, expected_updated_at }, { abortSignal }, gate) => {
+      const notebook = await getNotebook({ projectRef, id }, abortSignal, authHeaders)
+
+      if (notebook.updated_at !== expected_updated_at) {
+        throw new NotebookToolError(
+          `Notebook "${id}" changed since expected_updated_at (${expected_updated_at}); it is now ${notebook.updated_at}. Call get_notebook again and reissue run_notebook against the current content.`,
+          { exposeToAssistant: true }
+        )
+      }
+
+      const queryCells = notebook.content.cells.filter(
+        (cell) => cell._tag === 'database_cell' || cell._tag === 'log_cell'
+      )
+      const shouldLookupReplica = queryCells.some(
+        (cell) =>
+          cell._tag === 'database_cell' &&
+          cell.database_identifier !== undefined &&
+          cell.database_identifier !== projectRef
+      )
+      let databases: Awaited<ReturnType<typeof getReadReplicas>> = []
+      let replicaLookupFailure: { error: unknown } | undefined
+      if (shouldLookupReplica) {
+        try {
+          databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders)
+        } catch (error) {
+          replicaLookupFailure = { error }
+        }
+      }
+
+      const cells: NotebookRunCellOutput[] = []
+
+      // This execute function only runs after the AI SDK approval gate resolves true. That
+      // approval is the explicit user gesture which promotes every persisted cell's
+      // untrusted SQL immediately before execution. Run sequentially to preserve notebook
+      // order and avoid racing cells that may depend on earlier writes.
+      for (const cell of queryCells) {
+        const title = cell.title?.trim() || 'Untitled query'
+
+        try {
+          if (cell._tag === 'log_cell') {
+            const result = await executeLogsSql({
+              projectRef: projectRef ?? '',
+              sql: gate.acceptUntrustedLogsSql(cell.unchecked_sql),
+              range: resolveLogTimeRange(cell.time_range),
+              endpoint: QUERY_SOURCE_REGISTRY.logs.endpoint,
+              signal: abortSignal,
+              headers: authHeaders,
+            })
+            if (result.error) throw result.error
+            cells.push({
+              cell_id: cell._id,
+              title,
+              source: 'logs',
+              status: 'success',
+              // The analytics response crossed a JSON fetch boundary, although its shared
+              // data-layer type remains `unknown[]` because individual endpoints vary.
+              rows: result.rows as JSONValue[],
+            })
+            continue
+          }
+
+          const isReplicaTarget =
+            cell.database_identifier !== undefined && cell.database_identifier !== projectRef
+
+          if (isReplicaTarget && replicaLookupFailure !== undefined) {
+            throw replicaLookupFailure.error
+          }
+
+          const selectedConnectionString = !isReplicaTarget
+            ? connectionString
+            : databases?.find((database) => database.identifier === cell.database_identifier)
+                ?.connectionString
+
+          if (!selectedConnectionString) {
+            throw new Error(
+              `Unable to run query: connection string is missing for database "${cell.database_identifier ?? projectRef ?? 'primary'}"`
+            )
+          }
+
+          const limitedSql = applyAutoLimit(
+            gate.acceptUntrustedSql(cell.unchecked_sql),
+            cell.row_limit
+          )
+          const { result } = await executeSql<JSONValue[]>(
+            {
+              projectRef,
+              connectionString: selectedConnectionString,
+              sql: limitedSql.sql,
+              isStatementTimeoutDisabled: true,
+            },
+            abortSignal,
+            authHeaders
+          )
+          cells.push({
+            cell_id: cell._id,
+            title,
+            source: 'database',
+            status: 'success',
+            rows: Array.isArray(result) ? result : [],
+          })
+        } catch (error) {
+          cells.push({
+            cell_id: cell._id,
+            title,
+            source: cell._tag === 'log_cell' ? 'logs' : 'database',
+            status: 'error',
+            error: { message: getErrorMessage(error) },
+          })
+        }
+      }
+
+      return {
+        id: notebook.id,
+        name: notebook.name,
+        updated_at: notebook.updated_at,
+        optInLevel: aiOptInLevel,
+        cells,
+      } satisfies NotebookRunOutput
+    }
+  )
+}

apps/studio/lib/ai/tools/run-notebook.ts

@@ -0,0 +1,33 @@
+import { z } from 'zod'
+
+import {
+  INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE,
+  notebookRunOutputSchema,
+  sanitizeNotebookRunOutput,
+  type NotebookRunOutput,
+} from './notebook-run-output'
+import type { ToolContract } from './types'
+
+const inputSchema = z.object({
+  id: z.string().describe('The id of the notebook to run.'),
+  expected_updated_at: z
+    .string()
+    .describe(
+      'The `updated_at` received from `get_notebook`. The run is rejected if the notebook changed since.'
+    ),
+})
+
+export const contract = {
+  name: 'run_notebook',
+  description:
+    'Asks the user to run every database and logs query cell in a notebook, in notebook order, and returns each result according to the organization data-sharing permissions. Requires one user approval before any cell executes.',
+  inputSchema,
+  outputSchema: notebookRunOutputSchema,
+  minOptInLevel: 'schema',
+  needsApproval: true,
+  toModelOutput: (output: NotebookRunOutput, optInLevel) =>
+    sanitizeNotebookRunOutput(output, optInLevel),
+  // Persisted output comes from the client, so a malformed run is replaced with an
+  // explanation rather than dropped.
+  outputOnInvalid: INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE,
+} as const satisfies ToolContract

apps/studio/lib/ai/tools/studio-tools.ts

@@ -9,6 +9,7 @@ import { DEFAULT_CELL_ROW_LIMIT } from '@/components/interfaces/Explorer/QueryCe
 import { deployEdgeFunction } from '@/data/edge-functions/edge-functions-deploy-mutation'
 import { executeSql } from '@/data/sql/execute-sql-mutation'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
+import { isOptInLevelAtLeast } from '@/lib/ai/opt-in-level'
 import {
   EDGE_FUNCTION_PROMPT,
   LOGS_PROMPT,
@@ -17,7 +18,7 @@ import {
   RLS_PROMPT,
   STORAGE_PROMPT,
 } from '@/lib/ai/prompts'
-import { isOptInLevelAtLeast, updateOptInLevelInputSchema } from '@/lib/ai/tool-filter'
+import { updateOptInLevelInputSchema } from '@/lib/ai/tool-filter'
 import { NO_DATA_PERMISSIONS, toModelRows } from '@/lib/ai/tools/tool-sanitizer'
 import { fixSqlBackslashEscapes } from '@/lib/ai/util'
 

apps/studio/lib/ai/tools/tool-sanitizer.ts

@@ -1,8 +1,8 @@
 import type { ToolUIPart, UIMessage } from 'ai'
 import { z } from 'zod'
 
-import { lowerOptInLevel, optInLevelSchema, type ToolName } from '../tool-filter'
-import { sanitizeNotebookRunOutput } from './notebook-run-output'
+import { lowerOptInLevel, optInLevelSchema } from '../opt-in-level'
+import type { ToolName } from '../tool-filter'
 import { sanitizeContractPart } from './registry'
 import { DEFAULT_CELL_ROW_LIMIT } from '@/components/interfaces/Explorer/QueryCell/QueryCell.utils'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
@@ -68,17 +68,8 @@ const executeSqlSanitizer: ToolSanitizer = {
   },
 }
 
-const runNotebookSanitizer: ToolSanitizer = {
-  toolName: 'run_notebook',
-  sanitize: (tool, optInLevel) => ({
-    ...tool,
-    output: sanitizeNotebookRunOutput(tool.output, optInLevel),
-  }),
-}
-
 export const ALL_TOOL_SANITIZERS = {
   [executeSqlSanitizer.toolName]: executeSqlSanitizer,
-  [runNotebookSanitizer.toolName]: runNotebookSanitizer,
 }
 
 export function sanitizeMessagePart(

apps/studio/lib/ai/tools/types.ts

@@ -27,6 +27,11 @@ export type ToolContract<TInput = unknown, TOutput = unknown> = {
    * its own result. Return the output unchanged only if all of it is safe at every level.
    */
   toModelOutput(output: TOutput, optInLevel: AiOptInLevel): JSONValue
+  /**
+   * What history replay shows the model when a persisted output doesn't match `outputSchema`.
+   * Omit it to drop the output.
+   */
+  outputOnInvalid?: JSONValue
 }
 
 export type ToolContractInput<C extends ToolContract> = z.infer<C['inputSchema']>

apps/studio/lib/ai/tools/update-notebook.ts

@@ -1,5 +1,6 @@
 import { z } from 'zod'
 
+import { notebookMutationOutputSchema } from './notebook-mutation-output'
 import type { ToolContract } from './types'
 import { notebookOperationsSchema } from '@/data/content/notebooks/notebook-operations'
 import { notebookSchema } from '@/data/content/notebooks/notebook-schema'
@@ -20,9 +21,7 @@ const inputSchema = z.object({
  * `previous_content` lets the client re-derive the diff it already showed for approval,
  * without re-fetching a notebook that's now post-update. It is UI-only: `toModelOutput` drops it.
  */
-const outputSchema = z.object({
-  id: z.string(),
-  name: z.string(),
+const outputSchema = notebookMutationOutputSchema.extend({
   previous_content: notebookSchema.optional(),
 })