refactor(studio): add tool contracts and migrate update_notebook (#51626)

## Summary

Each assistant tool's identity and contract is currently written in up
to nine places that nothing ties together at compile time. This PR adds
a contract/server split and migrates `update_notebook`, the tool with
the most copies, as the first proof of the design.

- **Contract** (`lib/ai/tools/update-notebook.ts`): isomorphic file with
the input and output schemas, approval flag, minimum opt-in level, and
`toModelOutput`, the single rule for what the model may see.
- **Server** (`update-notebook.server.ts`): the executable tool, built
with `approvalGatedToolFromContract`.
- **Registry** (`registry.ts`): collects contracts, derives a typed
`StudioUIMessage`, and provides `sanitizeContractPart` for history
replay.
- **Shared types** (`types.ts`): `ToolContract`, `ToolContext`,
`ToolFactory`, plus a single `getAuthHeaders` replacing three variants.
- **Approval gate** (`approval-gate.ts`): approval-gated tools get
`acceptUntrustedSql` / `acceptUntrustedLogsSql` through a gate passed to
`execute`, and `*.server.ts` files can't import the raw functions
(ESLint).
- **Opt-in level**: `update_notebook` is gated by its contract's
`minOptInLevel`. `TOOL_CATEGORY_MAP` is now typed against `ToolName` and
excludes migrated tools, which exposed `get_active_incidents` missing
from the name enum.

Removed duplicates: the `update_notebook` sanitizer in
`tool-sanitizer.ts`, the client-side `stripNotebookSnapshot`, and the
copied update schemas in `Message.utils.ts`.

## Behavior changes

- The client no longer strips `previous_content` from `update_notebook`
outputs before sending history, so requests carry the previous notebook
snapshot. The server still strips it before the model sees it.
- On history replay, an `update_notebook` output that doesn't match its
schema is dropped instead of passed through minus `previous_content`.

## Testing

- New `registry.test.ts` (typed parts, replay sanitizer) and a contract
case in `tool-filter.test.ts`.
- Studio typecheck, lint ratchet, knip, Prettier and the affected vitest
suites pass.

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

@@ -1,8 +1,7 @@
 import { untrustedSql } from '@supabase/pg-meta'
 import { z, type SafeParseReturnType } from 'zod'
 
-import { notebookOperationsSchema } from '@/data/content/notebooks/notebook-operations'
-import { agentNotebookSchema, notebookSchema } from '@/data/content/notebooks/notebook-schema'
+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.
@@ -135,12 +134,6 @@ export const createNotebookInputSchema = z.object({
   content: agentNotebookSchema,
 })
 
-export const updateNotebookInputSchema = z.object({
-  id: z.string(),
-  expected_updated_at: z.string(),
-  operations: notebookOperationsSchema,
-})
-
 export const deleteNotebookInputSchema = z.object({
   id: z.string(),
 })
@@ -152,10 +145,6 @@ export const runNotebookInputSchema = z.object({
 
 export const notebookToolOutputSchema = z.object({ id: z.string(), name: z.string() })
 
-export const updateNotebookToolOutputSchema = notebookToolOutputSchema.extend({
-  previous_content: notebookSchema.optional(),
-})
-
 export const rateMessageResponseSchema = z.object({
   category: z.enum([
     'sql_generation',

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

@@ -13,8 +13,6 @@ import {
   createNotebookInputSchema,
   deleteNotebookInputSchema,
   notebookToolOutputSchema,
-  updateNotebookInputSchema,
-  updateNotebookToolOutputSchema,
 } from './Message.utils'
 import { AlertError } from '@/components/ui/AlertError'
 import { hasDiscardableChanges } from '@/data/content/notebooks/notebook-cache'
@@ -25,6 +23,7 @@ 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 updateNotebookContract } from '@/lib/ai/tools/update-notebook'
 import { useNotebooksStateSnapshot } from '@/state/notebooks/notebooks-state'
 
 export type NotebookProposalMode = 'create' | 'update' | 'delete'
@@ -386,7 +385,7 @@ function UpdateNotebookProposal({
 }: NotebookProposalStepPropsWithOutput) {
   const { ref } = useParams()
   const notebooksSnap = useNotebooksStateSnapshot()
-  const parsedInput = updateNotebookInputSchema.safeParse(input)
+  const parsedInput = updateNotebookContract.inputSchema.safeParse(input)
   const isCompleted = state === 'output-available'
   const hasUnsavedLocalChanges =
     confirmState === 'approval-requested' &&
@@ -415,7 +414,7 @@ function UpdateNotebookProposal({
   }
 
   if (isCompleted) {
-    const parsedOutput = updateNotebookToolOutputSchema.safeParse(output)
+    const parsedOutput = updateNotebookContract.outputSchema.safeParse(output)
     const notebookName = parsedOutput.success ? parsedOutput.data.name : undefined
     const isOutputForRequestedNotebook =
       parsedOutput.success && parsedOutput.data.id === parsedInput.data.id

apps/studio/eslint.config.cjs

@@ -39,6 +39,37 @@ const NO_MODULE_SCOPE_CREATE_CLIENT = {
     'Do not construct a Supabase client at module scope in API route files — the TanStack server evaluates every route module at boot, so a missing env var (e.g. SUPABASE_URL on platform) crashes every route. Construct it lazily inside the handler (see lib/api/self-hosted-admin.ts).',
 }
 
+const LUCIDE_ICON_IMPORT_PATTERNS = [
+  {
+    group: ['lucide-react'],
+    importNamePattern: '^(?!.*Icon$)(?!Lucide)(?!icons$)(?!create)',
+    message:
+      'Import Lucide icons with the Icon suffix (e.g. `CheckIcon`, not `Check` or `Check as CheckIcon`).',
+  },
+]
+
+const BASE_RESTRICTED_IMPORT_PATHS = [
+  {
+    name: 'components/ui/DataTable/DataTableColumn/DataTableColumnHeader',
+    message: 'Use TanStackTableHeadSort from ui-patterns/Table instead.',
+  },
+]
+
+const APPROVAL_GATE_MESSAGE =
+  'Promote untrusted SQL through the `gate` that approvalGatedToolFromContract passes to execute.'
+const TOOL_MODULE_RESTRICTED_IMPORT_PATHS = [
+  {
+    name: '@supabase/pg-meta',
+    importNames: ['acceptUntrustedSql'],
+    message: APPROVAL_GATE_MESSAGE,
+  },
+  {
+    name: '@/data/logs/safe-analytics-sql',
+    importNames: ['acceptUntrustedLogsSql'],
+    message: APPROVAL_GATE_MESSAGE,
+  },
+]
+
 module.exports = defineConfig([
   { files: ['**/*.ts', '**/*.tsx'] },
   supabaseConfig,
@@ -62,20 +93,8 @@ module.exports = defineConfig([
       'no-restricted-imports': [
         'error',
         {
-          paths: [
-            {
-              name: 'components/ui/DataTable/DataTableColumn/DataTableColumnHeader',
-              message: 'Use TanStackTableHeadSort from ui-patterns/Table instead.',
-            },
-          ],
-          patterns: [
-            {
-              group: ['lucide-react'],
-              importNamePattern: '^(?!.*Icon$)(?!Lucide)(?!icons$)(?!create)',
-              message:
-                'Import Lucide icons with the Icon suffix (e.g. `CheckIcon`, not `Check` or `Check as CheckIcon`).',
-            },
-          ],
+          paths: BASE_RESTRICTED_IMPORT_PATHS,
+          patterns: LUCIDE_ICON_IMPORT_PATTERNS,
         },
       ],
       'barrel-files/avoid-re-export-all': 'error',
@@ -119,6 +138,18 @@ module.exports = defineConfig([
       'no-restricted-syntax': ['error', ...ANALYTICS_SQL_RESTRICTED_SYNTAX],
     },
   },
+  {
+    files: ['lib/ai/tools/*.server.ts'],
+    rules: {
+      'no-restricted-imports': [
+        'error',
+        {
+          paths: [...BASE_RESTRICTED_IMPORT_PATHS, ...TOOL_MODULE_RESTRICTED_IMPORT_PATHS],
+          patterns: LUCIDE_ICON_IMPORT_PATTERNS,
+        },
+      ],
+    },
+  },
   // API route modules are eagerly imported by the TanStack server at boot, so
   // module-scope side effects there are especially dangerous. This block also
   // re-includes the analytics selectors because flat config replaces (not

apps/studio/lib/ai/generate-assistant-response.utils.test.ts

@@ -2,7 +2,7 @@ import type { ToolUIPart, UIMessage } from 'ai'
 import { describe, expect, it } from 'vitest'
 
 import { prepareMessagesForModel } from './generate-assistant-response.utils'
-import { encodeNotebookToolError, NotebookToolError } from './tools/notebook-tools'
+import { encodeNotebookToolError, NotebookToolError } from './tools/notebook-tool-error'
 
 function assistantMessage(parts: UIMessage['parts']): UIMessage {
   return { id: 'msg-1', role: 'assistant', parts }

apps/studio/lib/ai/generate-assistant-response.utils.ts

@@ -1,7 +1,7 @@
 import { isToolUIPart, type DynamicToolUIPart, type ToolUIPart, type UIMessage } from 'ai'
 
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
-import { decodeNotebookToolError } from '@/lib/ai/tools/notebook-tools'
+import { decodeNotebookToolError } from '@/lib/ai/tools/notebook-tool-error'
 import { sanitizeMessagePart } from '@/lib/ai/tools/tool-sanitizer'
 
 const INVALID_TOOL_STATES = [

apps/studio/lib/ai/message-utils.test.ts

@@ -1,4 +1,4 @@
-import type { DynamicToolUIPart, ToolUIPart, UIMessage } from 'ai'
+import type { DynamicToolUIPart, UIMessage } from 'ai'
 import { describe, expect, it } from 'vitest'
 
 import {
@@ -6,7 +6,6 @@ import {
   isManualApprovalRequested,
   prepareMessagesForAPI,
 } from './message-utils'
-import { createAssistantMessageWithUpdateNotebookTool } from './test-fixtures'
 
 const makeApprovalPart = (id: string, isAutomatic = false): DynamicToolUIPart =>
   ({
@@ -263,38 +262,4 @@ describe('prepareMessagesForAPI', () => {
     expect(result[2]).toEqual(messages[2])
     expect(result[3]).not.toHaveProperty('results')
   })
-
-  it('strips update_notebook previous_content before re-uploading to the API', () => {
-    const messages = [createAssistantMessageWithUpdateNotebookTool()]
-
-    const result = prepareMessagesForAPI(messages)
-
-    expect((result[0].parts[0] as ToolUIPart).output).toEqual({
-      id: 'notebook-1',
-      name: 'Signup funnel',
-    })
-  })
-
-  it('does not mutate the original message parts when stripping previous_content', () => {
-    const messages = [createAssistantMessageWithUpdateNotebookTool()]
-    const originalParts = messages[0].parts
-
-    prepareMessagesForAPI(messages)
-
-    expect(messages[0].parts).toBe(originalParts)
-    expect((originalParts[0] as ToolUIPart).output).toHaveProperty('previous_content')
-  })
-
-  it('leaves an update_notebook output without previous_content unchanged', () => {
-    const messages = [
-      createAssistantMessageWithUpdateNotebookTool({ id: 'notebook-1', name: 'Signup funnel' }),
-    ]
-
-    const result = prepareMessagesForAPI(messages)
-
-    expect((result[0].parts[0] as ToolUIPart).output).toEqual({
-      id: 'notebook-1',
-      name: 'Signup funnel',
-    })
-  })
 })

apps/studio/lib/ai/message-utils.ts

@@ -8,15 +8,6 @@ import {
 
 type UIPart = UIMessagePart<UIDataTypes, UITools>
 
-/** Strips `update_notebook`'s `previous_content` snapshot — display-only, never re-uploaded. */
-function stripNotebookSnapshot(part: UIPart): UIPart {
-  if (!isToolUIPart(part) || part.type !== 'tool-update_notebook') return part
-  if (!part.output || typeof part.output !== 'object') return part
-
-  const { previous_content, ...sanitizedOutput } = part.output as Record<string, unknown>
-  return { ...part, output: sanitizedOutput } as UIPart
-}
-
 /**
  * Prepares messages for API transmission by cleaning and limiting history
  */
@@ -35,11 +26,6 @@ export function prepareMessagesForAPI(messages: UIMessage[]): UIMessage[] {
     if (message.role === 'assistant' && message.results) {
       delete cleanedMessage.results
     }
-    // Map into a new array rather than mutating in place — `parts` is shared by reference
-    // with the locally persisted message the assistant panel reads from.
-    if (cleanedMessage.parts) {
-      cleanedMessage.parts = cleanedMessage.parts.map(stripNotebookSnapshot)
-    }
     return cleanedMessage as UIMessage
   })
 

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

@@ -4,17 +4,22 @@ 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 { notebooksState } from '@/state/notebooks/notebooks-state'
 
 export type NotebookCacheEffect =
   | { _tag: 'upserted'; toolCallId: string; id: string }
   | { _tag: 'deleted'; toolCallId: string; id: string }
 
-const NOTEBOOK_MUTATION_TOOL_TYPES = new Set([
-  'tool-create_notebook',
-  'tool-update_notebook',
-  'tool-delete_notebook',
-])
+const NOTEBOOK_MUTATION_TOOLS = [
+  'create_notebook',
+  'update_notebook',
+  'delete_notebook',
+] as const satisfies readonly ToolName[]
+
+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 {
   return NOTEBOOK_MUTATION_TOOL_TYPES.has(part.type)

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

@@ -9,6 +9,7 @@ import {
   TOOL_CATEGORY_MAP,
   toolSetValidationSchema,
 } from './tool-filter'
+import { contract as updateNotebookContract } from './tools/update-notebook'
 
 describe('TOOL_CATEGORY_MAP', () => {
   it('should categorize tools correctly', () => {
@@ -18,6 +19,23 @@ describe('TOOL_CATEGORY_MAP', () => {
   })
 })
 
+describe('tools with a contract', () => {
+  const tools = {
+    update_notebook: { description: 'Update a notebook', execute: vitest.fn() },
+  } as unknown as ToolSet
+
+  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)
+
+    const blocked = filterToolsByOptInLevel(tools, 'disabled').update_notebook
+    const allowed = filterToolsByOptInLevel(tools, 'schema').update_notebook
+
+    expect(blocked.description).toContain('Requires opting in')
+    expect(allowed.description).toBe('Update a notebook')
+  })
+})
+
 describe('tool allowance by opt-in level', () => {
   // Helper function to get allowed tools for a given opt-in level
   function getAllowedTools(optInLevel: string) {

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

@@ -1,6 +1,7 @@
 import type { Tool, ToolSet } from 'ai'
 import { z } from 'zod'
 
+import { getToolContract, type ContractToolName } from './tools/registry'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
 
 // Add the DatabaseExtension type import
@@ -27,6 +28,7 @@ export const toolSetValidationSchema = z.record(
     'search_docs',
     'get_advisors',
     'query_logs',
+    'get_active_incidents',
 
     // Local tools
     'execute_sql',
@@ -107,9 +109,10 @@ export const TOOL_CATEGORIES = {
 type ToolCategory = (typeof TOOL_CATEGORIES)[keyof typeof TOOL_CATEGORIES]
 
 /**
- * Mapping of tool names to their categories
+ * Mapping of tool names to their categories. Tools that have migrated to a contract
+ * (`TOOL_CONTRACTS`) declare their own `minOptInLevel` instead and aren't listed here.
  */
-export const TOOL_CATEGORY_MAP: Record<string, ToolCategory> = {
+export const TOOL_CATEGORY_MAP: Record<Exclude<ToolName, ContractToolName>, ToolCategory> = {
   // UI tools - always available
   execute_sql: TOOL_CATEGORIES.UI,
   deploy_edge_function: TOOL_CATEGORIES.UI,
@@ -134,7 +137,6 @@ export const TOOL_CATEGORY_MAP: Record<string, ToolCategory> = {
   get_notebook: TOOL_CATEGORIES.SCHEMA,
   run_notebook: TOOL_CATEGORIES.SCHEMA,
   create_notebook: TOOL_CATEGORIES.SCHEMA,
-  update_notebook: TOOL_CATEGORIES.SCHEMA,
   delete_notebook: TOOL_CATEGORIES.SCHEMA,
   getSchemaTables: TOOL_CATEGORIES.SCHEMA,
   getRlsKnowledge: TOOL_CATEGORIES.SCHEMA,
@@ -164,18 +166,31 @@ function getMinimumOptInLevel(category: ToolCategory): AiOptInLevel | null {
   }
 }
 
+function isCategorizedTool(name: string): name is keyof typeof TOOL_CATEGORY_MAP {
+  return Object.hasOwn(TOOL_CATEGORY_MAP, name)
+}
+
+/**
+ * The lowest opt-in level a tool needs, `null` if it's always available, or `undefined` if the
+ * tool isn't on the allowlist. A contract's `minOptInLevel` wins over the category map.
+ */
+function getToolMinimumLevel(toolName: string): AiOptInLevel | null | undefined {
+  const contract = getToolContract(toolName)
+  if (contract) return contract.minOptInLevel
+  if (isCategorizedTool(toolName)) return getMinimumOptInLevel(TOOL_CATEGORY_MAP[toolName])
+  return undefined
+}
+
 /**
  * Check if a tool is allowed based on the current opt-in level
  */
 function isToolAllowed(toolName: string, aiOptInLevel: AiOptInLevel): boolean {
-  const category = TOOL_CATEGORY_MAP[toolName]
+  const minimumLevel = getToolMinimumLevel(toolName)
 
-  if (!category) {
+  if (minimumLevel === undefined) {
     return false
   }
 
-  const minimumLevel = getMinimumOptInLevel(category)
-
   // UI tools are always allowed
   if (!minimumLevel) {
     return true
@@ -218,7 +233,7 @@ export function createPrivacyMessageTool(toolInstance: Tool<any, any>) {
 export function filterToolsByOptInLevel(tools: ToolSet, aiOptInLevel: AiOptInLevel): ToolSet {
   return Object.fromEntries(
     Object.entries(tools)
-      .filter(([toolName]) => TOOL_CATEGORY_MAP[toolName] !== undefined)
+      .filter(([toolName]) => getToolMinimumLevel(toolName) !== undefined)
       .map(([toolName, toolInstance]) => {
         if (isToolAllowed(toolName, aiOptInLevel)) {
           return [toolName, toolInstance]

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

@@ -0,0 +1,18 @@
+import { acceptUntrustedSql } from '@supabase/pg-meta'
+
+import { acceptUntrustedLogsSql } from '@/data/logs/safe-analytics-sql'
+
+/**
+ * What an approval-gated tool's `execute` uses to turn untrusted SQL into trusted SQL.
+ *
+ * Promoting untrusted SQL is only valid once the user has approved the tool call. The SDK
+ * calls `execute` only after approval, and `approvalGatedToolFromContract` is the only place
+ * that hands this gate to an `execute`. `*.server.ts` tool files can't import the raw
+ * `accept*` functions (see `eslint.config.cjs`), so they can't promote SQL without it.
+ */
+export type ApprovalGate = {
+  acceptUntrustedSql: typeof acceptUntrustedSql
+  acceptUntrustedLogsSql: typeof acceptUntrustedLogsSql
+}
+
+export const approvalGate: ApprovalGate = { acceptUntrustedSql, acceptUntrustedLogsSql }

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

@@ -0,0 +1,14 @@
+/**
+ * Headers for forwarding the caller's credentials to the platform API.
+ * Returns `undefined` when there's nothing to forward.
+ */
+export function getAuthHeaders(
+  authorization: string | undefined,
+  { contentType = false }: { contentType?: boolean } = {}
+): Record<string, string> | undefined {
+  if (!authorization) return undefined
+  return {
+    ...(contentType && { 'Content-Type': 'application/json' }),
+    Authorization: authorization,
+  }
+}

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

@@ -4,6 +4,8 @@ import { IS_PLATFORM } from 'common'
 import { stripIndent } from 'common-tags'
 import { z } from 'zod'
 
+import { getAuthHeaders } from './auth-headers'
+import type { ToolContext } from './types'
 import { getDatabaseFunctions } from '@/data/database-functions/database-functions-query'
 import { getDatabasePolicies } from '@/data/database-policies/database-policies-query'
 import { executeSql } from '@/data/sql/execute-sql-mutation'
@@ -15,17 +17,16 @@ export const getFallbackTools = ({
   cookie,
   authorization,
   includeSchemaMetadata,
-}: {
+}: ToolContext & {
   projectRef: string
   connectionString: string
   cookie?: string
-  authorization?: string
   includeSchemaMetadata: boolean
 }) => {
   const headers = {
     'Content-Type': 'application/json',
     ...(cookie && { cookie }),
-    ...(authorization && { Authorization: authorization }),
+    ...getAuthHeaders(authorization),
   }
 
   return {

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

@@ -0,0 +1,71 @@
+import { z } from 'zod'
+
+const notebookToolErrorMetadataSchema = z.object({
+  exposeToAssistant: z.boolean(),
+})
+
+export type NotebookToolErrorMetadata = z.infer<typeof notebookToolErrorMetadataSchema>
+
+const notebookToolErrorSchema = notebookToolErrorMetadataSchema.extend({
+  tag: z.literal('notebook_tool_error'),
+  message: z.string(),
+})
+
+export type EncodedNotebookToolError = z.infer<typeof notebookToolErrorSchema>
+
+/** Thrown by update_notebook for failures the assistant can act on by retrying. */
+export class NotebookToolError extends Error {
+  readonly metadata: NotebookToolErrorMetadata
+
+  constructor(message: string, metadata: NotebookToolErrorMetadata) {
+    super(message)
+    this.name = 'NotebookToolError'
+    this.metadata = notebookToolErrorMetadataSchema.parse(metadata)
+  }
+}
+
+/** Called from the stream's `onError` (generate-v4.ts), which still has the live Error. */
+export function encodeNotebookToolError(error: unknown): string | null {
+  if (!(error instanceof NotebookToolError)) return null
+  return JSON.stringify(
+    notebookToolErrorSchema.parse({
+      ...error.metadata,
+      tag: 'notebook_tool_error',
+      message: error.message,
+    })
+  )
+}
+
+export function decodeNotebookToolError(errorText: string): EncodedNotebookToolError | null {
+  let parsed: unknown
+  try {
+    parsed = JSON.parse(errorText)
+  } catch {
+    return null
+  }
+  const result = notebookToolErrorSchema.safeParse(parsed)
+  return result.success ? result.data : null
+}
+
+/**
+ * Rejects a made-up `database_identifier` before it's written: it would pass schema
+ * validation (it's just a string) but silently break the cell at run time, since
+ * QueryEditor's connection-string lookup can't resolve an identifier that isn't real.
+ *
+ * @throws NotebookToolError if any cell has a `database_identifier` that isn't in the
+ * list of valid identifiers for this project.
+ */
+export function assertValidDatabaseIdentifiers(
+  cells: ReadonlyArray<{ _tag: string; database_identifier?: string }>,
+  validIdentifiers: Set<string>
+): void {
+  for (const cell of cells) {
+    if (cell._tag !== 'database_cell' || cell.database_identifier === undefined) continue
+    if (!validIdentifiers.has(cell.database_identifier)) {
+      throw new NotebookToolError(
+        `Unknown database_identifier "${cell.database_identifier}" — call list_databases to see this project's valid identifiers.`,
+        { exposeToAssistant: true }
+      )
+    }
+  }
+}

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

@@ -6,9 +6,9 @@ import { sanitizeNotebookRunOutput, type NotebookRunOutput } from './notebook-ru
 import {
   decodeNotebookToolError,
   encodeNotebookToolError,
-  getNotebookTools,
   NotebookToolError,
-} from './notebook-tools'
+} from './notebook-tool-error'
+import { getNotebookTools } from './notebook-tools'
 import type { AgentNotebook } from '@/data/content/notebooks/notebook-schema'
 import { addAPIMock, type APIErrorBody } from '@/tests/lib/msw'
 

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

@@ -2,19 +2,18 @@ import { acceptUntrustedSql, untrustedSql } from '@supabase/pg-meta'
 import { tool, type JSONValue } 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 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 {
-  applyNotebookOperations,
-  describeNotebookOperationError,
-  notebookOperationsSchema,
-} from '@/data/content/notebooks/notebook-operations'
 import { getNotebook } from '@/data/content/notebooks/notebook-query'
 import {
   agentNotebookSchema,
@@ -22,23 +21,15 @@ import {
   type WritableCell,
   type WritableNotebook,
 } from '@/data/content/notebooks/notebook-schema'
-import { createNotebook, upsertNotebook } from '@/data/content/notebooks/notebook-upsert-mutation'
+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 { getReadReplicas } from '@/data/read-replicas/replicas-query'
 import { executeSql } from '@/data/sql/execute-sql-mutation'
 import { applyAutoLimit } from '@/data/sql/utils'
-import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
 import type { Notebooks } from '@/types'
 
-export type NotebookToolsContext = {
-  projectRef?: string
-  connectionString?: string
-  authorization?: string
-  aiOptInLevel?: AiOptInLevel
-}
-
 function getErrorMessage(error: unknown): string {
   if (error instanceof Error) return error.message
   if (typeof error === 'object' && error !== null && 'message' in error) {
@@ -48,80 +39,9 @@ function getErrorMessage(error: unknown): string {
   return 'An unexpected error occurred while running this query.'
 }
 
-const notebookToolErrorMetadataSchema = z.object({
-  exposeToAssistant: z.boolean(),
-})
-
-export type NotebookToolErrorMetadata = z.infer<typeof notebookToolErrorMetadataSchema>
-
-const notebookToolErrorSchema = notebookToolErrorMetadataSchema.extend({
-  tag: z.literal('notebook_tool_error'),
-  message: z.string(),
-})
-
-export type EncodedNotebookToolError = z.infer<typeof notebookToolErrorSchema>
-
-/** Thrown by update_notebook for failures the assistant can act on by retrying. */
-export class NotebookToolError extends Error {
-  readonly metadata: NotebookToolErrorMetadata
-
-  constructor(message: string, metadata: NotebookToolErrorMetadata) {
-    super(message)
-    this.name = 'NotebookToolError'
-    this.metadata = notebookToolErrorMetadataSchema.parse(metadata)
-  }
-}
-
-/** Called from the stream's `onError` (generate-v4.ts), which still has the live Error. */
-export function encodeNotebookToolError(error: unknown): string | null {
-  if (!(error instanceof NotebookToolError)) return null
-  return JSON.stringify(
-    notebookToolErrorSchema.parse({
-      ...error.metadata,
-      tag: 'notebook_tool_error',
-      message: error.message,
-    })
-  )
-}
-
-/** Called from the history filter (generate-assistant-response.ts) on a persisted errorText. */
-export function decodeNotebookToolError(errorText: string): EncodedNotebookToolError | null {
-  let parsed: unknown
-  try {
-    parsed = JSON.parse(errorText)
-  } catch {
-    return null
-  }
-  const result = notebookToolErrorSchema.safeParse(parsed)
-  return result.success ? result.data : null
-}
-
-/**
- * Rejects a made-up `database_identifier` before it's written: it would pass schema
- * validation (it's just a string) but silently break the cell at run time, since
- * QueryEditor's connection-string lookup can't resolve an identifier that isn't real.
- *
- * @throws NotebookToolError if any cell has a `database_identifier` that isn't in the
- * list of valid identifiers for this project.
- */
-function assertValidDatabaseIdentifiers(
-  cells: ReadonlyArray<{ _tag: string; database_identifier?: string }>,
-  validIdentifiers: Set<string>
-): void {
-  for (const cell of cells) {
-    if (cell._tag !== 'database_cell' || cell.database_identifier === undefined) continue
-    if (!validIdentifiers.has(cell.database_identifier)) {
-      throw new NotebookToolError(
-        `Unknown database_identifier "${cell.database_identifier}" — call list_databases to see this project's valid identifiers.`,
-        { exposeToAssistant: true }
-      )
-    }
-  }
-}
-
-export const getNotebookTools = (ctx: NotebookToolsContext = {}) => {
+export const getNotebookTools = (ctx: ToolContext = {}) => {
   const { projectRef, connectionString, authorization, aiOptInLevel = 'schema' } = ctx
-  const authHeaders = authorization ? { Authorization: authorization } : undefined
+  const authHeaders = getAuthHeaders(authorization)
 
   return {
     list_databases: tool({
@@ -394,95 +314,7 @@ export const getNotebookTools = (ctx: NotebookToolsContext = {}) => {
         return { id: result.id, name }
       },
     }),
-    update_notebook: tool({
-      description:
-        'Asks the user to apply an ordered list of cell operations (insert, replace, delete, move) to an existing notebook. Requires user approval before updating. Re-fetches the notebook right before applying the operations and rejects the update if it changed since expected_updated_at.',
-      inputSchema: z.object({
-        id: z.string().describe('The id of the notebook to update.'),
-        expected_updated_at: z
-          .string()
-          .describe(
-            'The `updated_at` you received from `get_notebook`. The update is rejected if the notebook changed since.'
-          ),
-        operations: notebookOperationsSchema.describe(
-          'An ordered list of operations to apply to the notebook, addressing existing cells by id.'
-        ),
-      }),
-      needsApproval: true,
-      execute: async ({ id, expected_updated_at, operations }, { abortSignal }) => {
-        const newCells = operations.flatMap((operation) =>
-          operation._tag === 'insert_cell' || operation._tag === 'replace_cell'
-            ? [operation.cell]
-            : []
-        )
-        if (
-          newCells.some(
-            (cell) => cell._tag === 'database_cell' && cell.database_identifier !== undefined
-          )
-        ) {
-          const databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders)
-          assertValidDatabaseIdentifiers(
-            newCells,
-            new Set((databases ?? []).map((database) => database.identifier))
-          )
-        }
-
-        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 update_notebook against the current content.`,
-            { exposeToAssistant: true }
-          )
-        }
-
-        // The result of applyNotebookOperations is never written or executed as-is — every
-        // cell is re-promoted via acceptUntrustedSql/acceptUntrustedLogsSql further down in
-        // this same execute, right before the PUT.
-        const wireNotebook = toWireNotebook(notebook.content)
-
-        const result = applyNotebookOperations(wireNotebook, operations)
-        if (!result.success) {
-          throw new NotebookToolError(describeNotebookOperationError(result.error), {
-            exposeToAssistant: true,
-          })
-        }
-
-        // Same promotion as create_notebook above, inlined here for the same auditability
-        // reason: it must stay visible next to this tool's own `needsApproval: true`.
-        const cells: WritableNotebook['cells'] = result.notebook.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)) }
-          }
-        })
-
-        await upsertNotebook(
-          {
-            projectRef: projectRef ?? '',
-            id,
-            name: notebook.name,
-            description: notebook.description ?? undefined,
-            content: { schema_version: result.notebook.schema_version, cells },
-          },
-          abortSignal,
-          authHeaders
-        )
-
-        // `previous_content` lets the client re-derive the diff it already showed for
-        // approval, without re-fetching a notebook that's now post-update. Never reaches
-        // the model — see toModelOutput below.
-        return { id, name: notebook.name, previous_content: wireNotebook }
-      },
-      toModelOutput: ({ output }) => ({
-        type: 'json',
-        value: { id: output.id, name: output.name },
-      }),
-    }),
+    update_notebook: createUpdateNotebookTool(ctx),
     delete_notebook: tool({
       description: 'Asks the user to delete a notebook. Requires user approval before deleting.',
       inputSchema: z.object({

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

@@ -0,0 +1,48 @@
+import type { ToolUIPart } from 'ai'
+import { describe, expect, expectTypeOf, it } from 'vitest'
+
+import { sanitizeContractPart, type StudioUIMessage } from './registry'
+
+describe('TOOL_CONTRACTS', () => {
+  it('types contract tool parts on StudioUIMessage', () => {
+    type Part = Extract<
+      StudioUIMessage['parts'][number],
+      { type: 'tool-update_notebook'; state: 'output-available' }
+    >
+    expectTypeOf<Part['input']>().toHaveProperty('expected_updated_at')
+    expectTypeOf<Part['output']>().toHaveProperty('name')
+  })
+})
+
+describe('sanitizeContractPart', () => {
+  const part = {
+    type: 'tool-update_notebook',
+    toolCallId: 'call-1',
+    state: 'output-available',
+    input: {},
+    output: {
+      id: 'notebook-1',
+      name: 'Signup funnel',
+      previous_content: { schema_version: 1, cells: [] },
+    },
+  } as ToolUIPart
+
+  it('drops previous_content from the model output on history replay', () => {
+    expect(sanitizeContractPart(part, 'schema')?.output).toEqual({
+      id: 'notebook-1',
+      name: 'Signup funnel',
+    })
+  })
+
+  it('drops an output that does not match the schema', () => {
+    expect(sanitizeContractPart({ ...part, output: 'oops' } as ToolUIPart, 'schema')?.output).toBe(
+      undefined
+    )
+  })
+
+  it('returns undefined for tools without a contract', () => {
+    expect(
+      sanitizeContractPart({ ...part, type: 'tool-execute_sql' } as ToolUIPart, 'schema')
+    ).toBe(undefined)
+  })
+})

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

@@ -0,0 +1,58 @@
+import type { InferUITools, Tool, ToolUIPart, UIDataTypes, UIMessage } from 'ai'
+
+import type { ToolName } from '../tool-filter'
+import type { ToolContract, ToolContractInput, ToolContractOutput } from './types'
+import { contract as updateNotebook } from './update-notebook'
+import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
+
+/**
+ * Every tool that has migrated to a contract. Tools not listed here are still described
+ * by hand in the places the contract replaces.
+ */
+export const TOOL_CONTRACTS = {
+  update_notebook: updateNotebook,
+} as const satisfies Partial<Record<ToolName, ToolContract>>
+
+export type ContractToolName = keyof typeof TOOL_CONTRACTS
+
+type ContractTools = {
+  [K in ContractToolName]: Tool<
+    ToolContractInput<(typeof TOOL_CONTRACTS)[K]>,
+    ToolContractOutput<(typeof TOOL_CONTRACTS)[K]>
+  >
+}
+
+export type ContractUITools = InferUITools<ContractTools>
+
+/**
+ * A message whose `tool-<name>` parts are typed for every contract tool. The first two type
+ * parameters are the SDK defaults, since the assistant types neither message metadata nor
+ * custom data parts.
+ */
+export type StudioUIMessage = UIMessage<unknown, UIDataTypes, ContractUITools>
+
+/** The contract for a migrated tool, or `undefined` if the tool hasn't been migrated. */
+export function getToolContract(toolName: string): ToolContract | undefined {
+  return Object.hasOwn(TOOL_CONTRACTS, toolName)
+    ? TOOL_CONTRACTS[toolName as ContractToolName]
+    : undefined
+}
+
+/**
+ * Applies the contract's `toModelOutput` to a persisted tool part, for history replay (where
+ * the SDK doesn't call it). Returns `undefined` when the tool has no contract. An output that
+ * doesn't match the contract's schema is dropped rather than passed through unchecked.
+ */
+export function sanitizeContractPart<Part extends ToolUIPart>(
+  part: Part,
+  optInLevel: AiOptInLevel
+): Part | undefined {
+  const contract = getToolContract(part.type.slice('tool-'.length))
+  if (!contract) return undefined
+
+  const parsed = contract.outputSchema.safeParse(part.output)
+  return {
+    ...part,
+    output: parsed.success ? contract.toModelOutput(parsed.data, optInLevel) : undefined,
+  }
+}

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

@@ -1,18 +1,15 @@
 import { tool } from 'ai'
 import { z } from 'zod'
 
+import { getAuthHeaders } from './auth-headers'
+import type { ToolContext } from './types'
 import { getContentById } from '@/data/content/content-id-query'
 import { getContent } from '@/data/content/content-query'
 import type { Dashboards, SqlSnippets } from '@/types'
 
-export type ReportToolsContext = {
-  projectRef?: string
-  authorization?: string
-}
-
-export const getReportTools = (ctx: ReportToolsContext = {}) => {
+export const getReportTools = (ctx: ToolContext = {}) => {
   const { projectRef, authorization } = ctx
-  const authHeaders = authorization ? { Authorization: authorization } : undefined
+  const authHeaders = getAuthHeaders(authorization)
 
   return {
     list_reports: tool({

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

@@ -1,17 +1,15 @@
 import { tool } from 'ai'
 import { z } from 'zod'
 
+import { getAuthHeaders } from './auth-headers'
+import type { ToolContext } from './types'
 import { getDatabasePolicies } from '@/data/database-policies/database-policies-query'
 
 export const getSchemaTools = ({
   projectRef,
   connectionString,
   authorization,
-}: {
-  projectRef: string
-  connectionString: string
-  authorization?: string
-}) => ({
+}: ToolContext & { projectRef: string; connectionString: string }) => ({
   list_policies: tool({
     description: 'Get existing RLS policies for a given schema',
     inputSchema: z.object({
@@ -25,7 +23,7 @@ export const getSchemaTools = ({
           schemas,
         },
         abortSignal,
-        authorization ? { Authorization: authorization } : undefined
+        getAuthHeaders(authorization)
       )
 
       const formattedPolicies = data

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

@@ -2,6 +2,8 @@ import { acceptUntrustedSql, untrustedSql } from '@supabase/pg-meta'
 import { tool } from 'ai'
 import { z } from 'zod'
 
+import { getAuthHeaders } from './auth-headers'
+import type { ToolContext } from './types'
 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'
@@ -54,18 +56,9 @@ export const loadKnowledgeInputSchema = z.object({
     .describe('The knowledge to load'),
 })
 
-export type StudioToolsContext = {
-  projectRef?: string
-  connectionString?: string
-  authorization?: string
-  aiOptInLevel?: AiOptInLevel
-}
-
-export const getStudioTools = (ctx: StudioToolsContext = {}) => {
+export const getStudioTools = (ctx: ToolContext = {}) => {
   const { projectRef, connectionString, authorization, aiOptInLevel = 'schema' } = ctx
-  const authHeaders = authorization
-    ? { 'Content-Type': 'application/json', Authorization: authorization }
-    : undefined
+  const authHeaders = getAuthHeaders(authorization, { contentType: true })
 
   return {
     // Tools that wait on the user use `needsApproval` with a server-side `execute`, not a

apps/studio/lib/ai/tools/tool-from-contract.ts

@@ -0,0 +1,46 @@
+import type { Tool, ToolExecutionOptions } from 'ai'
+
+import { approvalGate, type ApprovalGate } from './approval-gate'
+import type { ToolContext, ToolContract } from './types'
+
+/**
+ * Builds the SDK tool for a contract. Description, input schema, approval and the model's
+ * view of the output all come from the contract; only `execute` is supplied.
+ */
+export function toolFromContract<TInput, TOutput>(
+  contract: ToolContract<TInput, TOutput>,
+  { aiOptInLevel = 'schema' }: ToolContext,
+  execute: (input: TInput, options: ToolExecutionOptions<unknown>) => PromiseLike<TOutput> | TOutput
+): Tool<TInput, TOutput> {
+  const { toModelOutput } = contract
+
+  // The SDK's `Tool` type is conditional on whether `TOutput` is `never`, which TypeScript
+  // can't resolve for a generic `TOutput`. Each factory's `ToolFactory<typeof contract>`
+  // annotation checks the concrete types instead.
+  return {
+    description: contract.description,
+    inputSchema: contract.inputSchema,
+    needsApproval: contract.needsApproval,
+    execute,
+    toModelOutput: ({ output }: { output: TOutput }) => ({
+      type: 'json' as const,
+      value: toModelOutput(output, aiOptInLevel),
+    }),
+  } as unknown as Tool<TInput, TOutput>
+}
+
+/**
+ * Like `toolFromContract`, for a contract that requires user approval. `execute` also receives
+ * the `ApprovalGate`, which is how it gets the functions that promote untrusted SQL.
+ */
+export function approvalGatedToolFromContract<TInput, TOutput>(
+  contract: ToolContract<TInput, TOutput> & { readonly needsApproval: true },
+  ctx: ToolContext,
+  execute: (
+    input: TInput,
+    options: ToolExecutionOptions<unknown>,
+    gate: ApprovalGate
+  ) => PromiseLike<TOutput> | TOutput
+): Tool<TInput, TOutput> {
+  return toolFromContract(contract, ctx, (input, options) => execute(input, options, approvalGate))
+}

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

@@ -3,6 +3,7 @@ import { z } from 'zod'
 
 import { lowerOptInLevel, optInLevelSchema, type ToolName } from '../tool-filter'
 import { sanitizeNotebookRunOutput } from './notebook-run-output'
+import { sanitizeContractPart } from './registry'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
 
 interface ToolSanitizer {
@@ -43,21 +44,6 @@ const executeSqlSanitizer: ToolSanitizer = {
   },
 }
 
-// `previous_content` is UI-only and must never reach the model — `toModelOutput` on the
-// tool covers the same turn, but history replay skips it, so it's stripped here too.
-const updateNotebookSanitizer: ToolSanitizer = {
-  toolName: 'update_notebook',
-  sanitize: (tool) => {
-    if (!tool.output || typeof tool.output !== 'object') return tool
-
-    const { previous_content, ...sanitizedOutput } = tool.output as Record<string, unknown>
-    return {
-      ...tool,
-      output: sanitizedOutput,
-    }
-  },
-}
-
 const runNotebookSanitizer: ToolSanitizer = {
   toolName: 'run_notebook',
   sanitize: (tool, optInLevel) => ({
@@ -68,7 +54,6 @@ const runNotebookSanitizer: ToolSanitizer = {
 
 export const ALL_TOOL_SANITIZERS = {
   [executeSqlSanitizer.toolName]: executeSqlSanitizer,
-  [updateNotebookSanitizer.toolName]: updateNotebookSanitizer,
   [runNotebookSanitizer.toolName]: runNotebookSanitizer,
 }
 
@@ -78,6 +63,8 @@ export function sanitizeMessagePart(
 ): UIMessage['parts'][number] {
   if (part.type.startsWith('tool-')) {
     const toolPart = part as ToolUIPart
+    const contractPart = sanitizeContractPart(toolPart, optInLevel)
+    if (contractPart) return contractPart
     const toolName = toolPart.type.slice('tool-'.length)
     const sanitizer = ALL_TOOL_SANITIZERS[toolName]
     if (sanitizer) {

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

@@ -0,0 +1,46 @@
+import type { JSONValue, Tool } from 'ai'
+import type { z } from 'zod'
+
+import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
+
+export type ToolContext = {
+  projectRef?: string
+  connectionString?: string
+  authorization?: string
+  aiOptInLevel?: AiOptInLevel
+}
+
+/**
+ * Isomorphic: zod and pure functions only, so nothing here may import the data layer.
+ */
+export type ToolContract<TInput = unknown, TOutput = unknown> = {
+  name: string
+  description: string
+  inputSchema: z.ZodType<TInput, z.ZodTypeDef, unknown>
+  outputSchema: z.ZodType<TOutput, z.ZodTypeDef, unknown>
+  /** Lowest opt-in level at which the tool runs. `null` means always available. */
+  minOptInLevel: AiOptInLevel | null
+  needsApproval: boolean
+  /**
+   * What the model may see of the output. Required, so a contract can't forget to decide.
+   * Applied to the current turn and again on history replay, so it must be safe to run on
+   * its own result. Return the output unchanged only if all of it is safe at every level.
+   */
+  toModelOutput(output: TOutput, optInLevel: AiOptInLevel): JSONValue
+}
+
+export type ToolContractInput<C extends ToolContract> = z.infer<C['inputSchema']>
+export type ToolContractOutput<C extends ToolContract> = z.infer<C['outputSchema']>
+
+/** The server half of a tool: builds the executable tool for a contract. */
+export type ToolFactory<C extends ToolContract> = (
+  ctx: ToolContext
+) => Tool<ToolContractInput<C>, ToolContractOutput<C>>
+
+/**
+ * A server half whose `execute` runs after the user approves. Only a contract declared with
+ * `needsApproval: true` satisfies the constraint. Build it with `approvalGatedToolFromContract`,
+ * which gives `execute` the `ApprovalGate` it needs to promote untrusted SQL.
+ */
+export type ApprovalGatedToolFactory<C extends ToolContract & { readonly needsApproval: true }> =
+  ToolFactory<C>

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

@@ -0,0 +1,97 @@
+import { untrustedSql } from '@supabase/pg-meta'
+
+import { getAuthHeaders } from './auth-headers'
+import { assertValidDatabaseIdentifiers, NotebookToolError } from './notebook-tool-error'
+import { approvalGatedToolFromContract } from './tool-from-contract'
+import type { ApprovalGatedToolFactory } from './types'
+import { contract } from './update-notebook'
+import {
+  applyNotebookOperations,
+  describeNotebookOperationError,
+} from '@/data/content/notebooks/notebook-operations'
+import { getNotebook } from '@/data/content/notebooks/notebook-query'
+import {
+  toWireNotebook,
+  type WritableCell,
+  type WritableNotebook,
+} from '@/data/content/notebooks/notebook-schema'
+import { upsertNotebook } 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 ({ id, expected_updated_at, operations }, { abortSignal }, gate) => {
+      const newCells = operations.flatMap((operation) =>
+        operation._tag === 'insert_cell' || operation._tag === 'replace_cell'
+          ? [operation.cell]
+          : []
+      )
+      if (
+        newCells.some(
+          (cell) => cell._tag === 'database_cell' && cell.database_identifier !== undefined
+        )
+      ) {
+        const databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders)
+        assertValidDatabaseIdentifiers(
+          newCells,
+          new Set((databases ?? []).map((database) => database.identifier))
+        )
+      }
+
+      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 update_notebook against the current content.`,
+          { exposeToAssistant: true }
+        )
+      }
+
+      // The result of applyNotebookOperations is never written or executed as-is — every
+      // cell is re-promoted via gate.acceptUntrustedSql/gate.acceptUntrustedLogsSql further
+      // down in this same execute, right before the PUT.
+      const wireNotebook = toWireNotebook(notebook.content)
+
+      const result = applyNotebookOperations(wireNotebook, operations)
+      if (!result.success) {
+        throw new NotebookToolError(describeNotebookOperationError(result.error), {
+          exposeToAssistant: true,
+        })
+      }
+
+      // Same promotion as create_notebook, inlined here for the same auditability reason: it
+      // must stay visible in the function the approval gate guards. `gate` is only handed to
+      // an `execute` that runs after the user approves.
+      const cells: WritableNotebook['cells'] = result.notebook.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)) }
+        }
+      })
+
+      await upsertNotebook(
+        {
+          projectRef: projectRef ?? '',
+          id,
+          name: notebook.name,
+          description: notebook.description ?? undefined,
+          content: { schema_version: result.notebook.schema_version, cells },
+        },
+        abortSignal,
+        authHeaders
+      )
+
+      return { id, name: notebook.name, previous_content: wireNotebook }
+    }
+  )
+}

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

@@ -0,0 +1,38 @@
+import { z } from 'zod'
+
+import type { ToolContract } from './types'
+import { notebookOperationsSchema } from '@/data/content/notebooks/notebook-operations'
+import { notebookSchema } from '@/data/content/notebooks/notebook-schema'
+
+const inputSchema = z.object({
+  id: z.string().describe('The id of the notebook to update.'),
+  expected_updated_at: z
+    .string()
+    .describe(
+      'The `updated_at` you received from `get_notebook`. The update is rejected if the notebook changed since.'
+    ),
+  operations: notebookOperationsSchema.describe(
+    'An ordered list of operations to apply to the notebook, addressing existing cells by id.'
+  ),
+})
+
+/**
+ * `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(),
+  previous_content: notebookSchema.optional(),
+})
+
+export const contract = {
+  name: 'update_notebook',
+  description:
+    'Asks the user to apply an ordered list of cell operations (insert, replace, delete, move) to an existing notebook. Requires user approval before updating. Re-fetches the notebook right before applying the operations and rejects the update if it changed since expected_updated_at.',
+  inputSchema,
+  outputSchema,
+  minOptInLevel: 'schema',
+  needsApproval: true,
+  toModelOutput: ({ id, name }) => ({ id, name }),
+} as const satisfies ToolContract

apps/studio/pages/api/ai/sql/generate-v4.ts

@@ -27,7 +27,7 @@ import {
   type AssistantModelId,
 } from '@/lib/ai/model.utils'
 import { getTools } from '@/lib/ai/tools'
-import { encodeNotebookToolError } from '@/lib/ai/tools/notebook-tools'
+import { encodeNotebookToolError } from '@/lib/ai/tools/notebook-tool-error'
 import { apiWrapper } from '@/lib/api/apiWrapper'
 import { executeQuery } from '@/lib/api/self-hosted/query'
 import { getBasePathURL } from '@/lib/helpers'