test(studio): mock notebook tools (#48952)

## Summary
- Adds deterministic `list_notebooks`/`get_notebook` fixtures (two
seeded notebooks, each with markdown/database/log cells) and stateful
in-memory `create_notebook`/`update_notebook` mocks to
`apps/studio/lib/ai/tools/mock-tools.ts`, so Braintrust evals can
exercise notebook tool calls without a real project.
- Both write-tool mocks force `needsApproval: false`, matching the
existing `execute_sql`/`deploy_edge_function` mock pattern — the eval
harness filters out tool-parts in `'approval-requested'` state and can
never answer an approval gate.
- All four notebook tools are wrapped from the real `getNotebookTools()`
definitions (only `execute`/`needsApproval` overridden), so evals
validate the model's arguments against the exact production schemas.
- Dedupes `describeOperationError` (previously duplicated between this
new mock and `notebook-tools.ts`) into a single exported
`describeNotebookOperationError` in `notebook-operations.ts`.

**Stacked on #48949** (`feature/notebooks-update-tool`) — this PR's base
branch is that PR, not `master`, because it reuses `update_notebook` and
the shared error helper that only exist there. Merge #48949 first, then
retarget/merge this one.

## Test plan
- [x] \`pnpm --filter studio typecheck\` passes
- [x] \`pnpm --filter studio test\` — all notebook-related suites pass
(\`mock-tools.test.ts\`, \`notebook-tools.test.ts\`,
\`data/content/notebooks/*\`)
- [x] \`eslint\` / \`prettier --check\` clean on all touched files

apps/studio/data/content/notebooks/notebook-operations.ts

@@ -67,6 +67,17 @@ export type ApplyNotebookOperationsResult =
   | { success: true; notebook: NotebookOperationsResult }
   | { success: false; error: NotebookOperationError }
 
+export function describeNotebookOperationError(error: NotebookOperationError): string {
+  switch (error._tag) {
+    case 'unknown_cell_id':
+      return `No cell with id "${error.cell_id}" exists in this notebook.`
+    case 'conflicting_operations':
+      return `More than one operation targets cell "${error.cell_id}".`
+    case 'empty_result':
+      return 'This update would leave the notebook with no cells.'
+  }
+}
+
 function targetCellId(operation: NotebookOperation): string | undefined {
   switch (operation._tag) {
     case 'insert_cell':

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

@@ -1,6 +1,8 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-import { getMockTools } from './mock-tools'
+import { getMockTools, MOCK_NOTEBOOKS_DATA } from './mock-tools'
+import { getNotebookTools } from './notebook-tools'
+import type { AgentNotebook } from '@/data/content/notebooks/notebook-schema'
 import { createInProcessSupabaseMCPClient } from '@/lib/ai/supabase-mcp'
 
 // The one real tool in the eval harness (search_docs) is sourced from an
@@ -61,4 +63,201 @@ describe('ai/tools/mock-tools getMockTools', () => {
   afterEach(() => {
     vi.clearAllMocks()
   })
+
+  describe('notebook tools', () => {
+    const AUTH_HEALTH_NOTEBOOK_ID = MOCK_NOTEBOOKS_DATA[0].id
+    const EDGE_FUNCTION_NOTEBOOK_ID = MOCK_NOTEBOOKS_DATA[1].id
+
+    it('list_notebooks reflects the two seeded fixtures', async () => {
+      const mockTools = await getMockTools(undefined, new AbortController().signal)
+      if (!mockTools.list_notebooks.execute) throw new Error('execute is undefined')
+
+      const result = await mockTools.list_notebooks.execute(
+        { limit: 20 },
+        { toolCallId: 'test', messages: [] }
+      )
+
+      expect(result.notebooks.map((notebook) => notebook.name)).toEqual([
+        'Auth health check',
+        'Edge function error triage',
+      ])
+      expect(result.notebooks.map((notebook) => notebook.cell_count)).toEqual([3, 2])
+      expect(result.notebooks[1].description).toBeUndefined()
+      expect(result.cursor).toBeUndefined()
+    })
+
+    it('get_notebook resolves cells in order and rejects an unknown id', async () => {
+      const mockTools = await getMockTools(undefined, new AbortController().signal)
+      if (!mockTools.get_notebook.execute) throw new Error('execute is undefined')
+
+      const result = await mockTools.get_notebook.execute(
+        { id: AUTH_HEALTH_NOTEBOOK_ID },
+        { toolCallId: 'test', messages: [] }
+      )
+
+      expect(result.cells.map((cell) => cell._tag)).toEqual([
+        'markdown_cell',
+        'database_cell',
+        'log_cell',
+      ])
+
+      const [, databaseCell, logCell] = result.cells
+      if (databaseCell._tag !== 'database_cell') throw new Error('expected database_cell')
+      if (logCell._tag !== 'log_cell') throw new Error('expected log_cell')
+
+      expect(databaseCell.sql).toContain('signups')
+      expect(databaseCell.row_limit).toBe(30)
+      expect(logCell.time_range).toEqual({ _tag: 'relative_time_range', unit: 'hour', amount: 1 })
+
+      await expect(
+        mockTools.get_notebook.execute(
+          { id: 'unknown-notebook-id' },
+          { toolCallId: 'test', messages: [] }
+        )
+      ).rejects.toThrow(/not found/i)
+    })
+
+    it('overrides create_notebook needsApproval to false, unlike the real tool', async () => {
+      const mockTools = await getMockTools(undefined, new AbortController().signal)
+
+      expect(getNotebookTools().create_notebook.needsApproval).toBe(true)
+      expect(mockTools.create_notebook.needsApproval).toBe(false)
+    })
+
+    it('create_notebook stores a new notebook visible via get_notebook and list_notebooks', async () => {
+      const mockTools = await getMockTools(undefined, new AbortController().signal)
+      if (!mockTools.create_notebook.execute) throw new Error('execute is undefined')
+      if (!mockTools.get_notebook.execute) throw new Error('execute is undefined')
+      if (!mockTools.list_notebooks.execute) throw new Error('execute is undefined')
+
+      const content: AgentNotebook = {
+        schema_version: 1,
+        cells: [
+          { _tag: 'markdown_cell', text: '# New notebook' },
+          { _tag: 'database_cell', sql: 'select 1', row_limit: 10 },
+        ],
+      }
+
+      const created = await mockTools.create_notebook.execute(
+        { name: 'New notebook', content },
+        { toolCallId: 'test', messages: [] }
+      )
+      expect(created).toEqual({ id: expect.any(String), name: 'New notebook' })
+
+      const fetched = await mockTools.get_notebook.execute(
+        { id: created.id },
+        { toolCallId: 'test', messages: [] }
+      )
+      expect(fetched.cells).toHaveLength(2)
+      expect(fetched.cells.every((cell) => typeof cell.id === 'string')).toBe(true)
+
+      const listed = await mockTools.list_notebooks.execute(
+        { limit: 20 },
+        { toolCallId: 'test', messages: [] }
+      )
+      expect(listed.notebooks).toHaveLength(3)
+      const newEntry = listed.notebooks.find((notebook) => notebook.id === created.id)
+      expect(newEntry?.cell_count).toBe(2)
+    })
+
+    it('overrides update_notebook needsApproval to false, unlike the real tool', async () => {
+      const mockTools = await getMockTools(undefined, new AbortController().signal)
+
+      expect(getNotebookTools().update_notebook.needsApproval).toBe(true)
+      expect(mockTools.update_notebook.needsApproval).toBe(false)
+    })
+
+    it('update_notebook inserts and deletes cells, and list_notebooks reflects the new cell count', async () => {
+      const mockTools = await getMockTools(undefined, new AbortController().signal)
+      if (!mockTools.get_notebook.execute) throw new Error('execute is undefined')
+      if (!mockTools.update_notebook.execute) throw new Error('execute is undefined')
+      if (!mockTools.list_notebooks.execute) throw new Error('execute is undefined')
+
+      const before = await mockTools.get_notebook.execute(
+        { id: AUTH_HEALTH_NOTEBOOK_ID },
+        { toolCallId: 'test', messages: [] }
+      )
+      const [markdownCell, , logCell] = before.cells
+
+      const result = await mockTools.update_notebook.execute(
+        {
+          id: AUTH_HEALTH_NOTEBOOK_ID,
+          operations: [
+            {
+              _tag: 'insert_cell',
+              after_cell_id: markdownCell.id,
+              cell: { _tag: 'database_cell', sql: 'select 1', row_limit: 10 },
+            },
+            { _tag: 'delete_cell', cell_id: logCell.id },
+          ],
+        },
+        { toolCallId: 'test', messages: [] }
+      )
+      expect(result).toEqual({ id: AUTH_HEALTH_NOTEBOOK_ID, name: 'Auth health check' })
+
+      const after = await mockTools.get_notebook.execute(
+        { id: AUTH_HEALTH_NOTEBOOK_ID },
+        { toolCallId: 'test', messages: [] }
+      )
+      expect(after.cells.map((cell) => cell._tag)).toEqual([
+        'markdown_cell',
+        'database_cell',
+        'database_cell',
+      ])
+
+      const listed = await mockTools.list_notebooks.execute(
+        { limit: 20 },
+        { toolCallId: 'test', messages: [] }
+      )
+      const entry = listed.notebooks.find((notebook) => notebook.id === AUTH_HEALTH_NOTEBOOK_ID)
+      expect(entry?.cell_count).toBe(3)
+    })
+
+    it('update_notebook rejects an unknown cell_id without mutating the notebook', async () => {
+      const mockTools = await getMockTools(undefined, new AbortController().signal)
+      if (!mockTools.get_notebook.execute) throw new Error('execute is undefined')
+      if (!mockTools.update_notebook.execute) throw new Error('execute is undefined')
+
+      await expect(
+        mockTools.update_notebook.execute(
+          {
+            id: EDGE_FUNCTION_NOTEBOOK_ID,
+            operations: [{ _tag: 'delete_cell', cell_id: 'does-not-exist' }],
+          },
+          { toolCallId: 'test', messages: [] }
+        )
+      ).rejects.toThrow(/does-not-exist/)
+
+      const after = await mockTools.get_notebook.execute(
+        { id: EDGE_FUNCTION_NOTEBOOK_ID },
+        { toolCallId: 'test', messages: [] }
+      )
+      expect(after.cells.map((cell) => cell._tag)).toEqual(['markdown_cell', 'log_cell'])
+    })
+
+    it('is isolated per call to getMockTools', async () => {
+      const firstCall = await getMockTools(undefined, new AbortController().signal)
+      if (!firstCall.create_notebook.execute) throw new Error('execute is undefined')
+
+      await firstCall.create_notebook.execute(
+        {
+          name: 'Ephemeral notebook',
+          content: { schema_version: 1, cells: [{ _tag: 'markdown_cell', text: 'hi' }] },
+        },
+        { toolCallId: 'test', messages: [] }
+      )
+
+      const secondCall = await getMockTools(undefined, new AbortController().signal)
+      if (!secondCall.list_notebooks.execute) throw new Error('execute is undefined')
+
+      const result = await secondCall.list_notebooks.execute(
+        { limit: 20 },
+        { toolCallId: 'test', messages: [] }
+      )
+      expect(result.notebooks.map((notebook) => notebook.name)).toEqual([
+        'Auth health check',
+        'Edge function error triage',
+      ])
+    })
+  })
 })

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

@@ -1,8 +1,20 @@
 import assert from 'node:assert'
-import { tool, type ToolSet } from 'ai'
+import { tool, type ToolCallOptions, type ToolSet } from 'ai'
 import { z } from 'zod'
 
 import { getStudioTools } from '../tools/studio-tools'
+import { getNotebookTools } from './notebook-tools'
+import {
+  applyNotebookOperations,
+  describeNotebookOperationError,
+  type NotebookOperation,
+  type OperationResultCell,
+} from '@/data/content/notebooks/notebook-operations'
+import type {
+  AgentNotebook,
+  CellWire,
+  NotebookWire,
+} from '@/data/content/notebooks/notebook-schema'
 import { createInProcessSupabaseMCPClient } from '@/lib/ai/supabase-mcp'
 
 const listTablesInputSchema = z.object({
@@ -143,6 +155,75 @@ const MOCK_LOGS_DATA = [
   },
 ]
 
+type MockNotebook = {
+  id: string
+  name: string
+  description?: string
+  visibility: 'project'
+  updated_at: string
+  content: NotebookWire
+}
+
+const MOCK_NOTEBOOK_TIMESTAMP = '2024-06-20T14:30:00Z'
+
+export const MOCK_NOTEBOOKS_DATA: MockNotebook[] = [
+  {
+    id: '6f1d3a54-8c2b-4d19-9f60-2a7b5c8e1d40',
+    name: 'Auth health check',
+    description: 'Daily signups plus any auth errors from the last hour.',
+    visibility: 'project',
+    updated_at: MOCK_NOTEBOOK_TIMESTAMP,
+    content: {
+      schema_version: 1,
+      cells: [
+        {
+          _tag: 'markdown_cell',
+          id: 'c1a0b8e2-3f47-4a52-9d18-6b0c4e2f7a91',
+          text: '# Auth health\n\nRun this daily: signup volume, then anything the auth service logged as an error.',
+        },
+        {
+          _tag: 'database_cell',
+          id: 'd2b1c9f3-4a58-4b63-8e29-7c1d5f3a8b02',
+          title: 'Signups per day',
+          sql: "select date_trunc('day', created_at) as day, count(*) as signups\nfrom auth.users\ngroup by day\norder by day desc",
+          row_limit: 30,
+          chart: { x_column: 'day', y_column: 'signups', cumulative: false },
+        },
+        {
+          _tag: 'log_cell',
+          id: 'e3c2d0a4-5b69-4c74-9f3a-8d2e6a4b9c13',
+          title: 'Auth errors',
+          sql: "select timestamp, event_message\nfrom auth_logs\nwhere event_message like '%error%'\norder by timestamp desc",
+          time_range: { _tag: 'relative_time_range', unit: 'hour', amount: 1 },
+        },
+      ],
+    },
+  },
+  {
+    id: '9a4e7b21-6d0c-4f38-8b57-3e1f9c6a2d84',
+    name: 'Edge function error triage',
+    visibility: 'project',
+    updated_at: MOCK_NOTEBOOK_TIMESTAMP,
+    content: {
+      schema_version: 1,
+      cells: [
+        {
+          _tag: 'markdown_cell',
+          id: 'f4d3e1b5-7c80-4d85-8a4b-9e3f7b5c0d24',
+          text: '# Edge function errors\n\nFailures from the last day, newest first.',
+        },
+        {
+          _tag: 'log_cell',
+          id: '0a5e4f2c-8d91-4e96-9b5c-af408c6d1e35',
+          title: 'hello-world failures',
+          sql: "select timestamp, event_message\nfrom function_edge_logs\nwhere event_message like '%TypeError%'\norder by timestamp desc",
+          time_range: { _tag: 'relative_time_range', unit: 'day', amount: 1 },
+        },
+      ],
+    },
+  },
+]
+
 function createMockedStudioTools() {
   const studioTools = getStudioTools()
 
@@ -267,6 +348,152 @@ function createMockListPoliciesTool() {
   })
 }
 
+function createMockNotebookStore() {
+  const notebooks = new Map(MOCK_NOTEBOOKS_DATA.map((notebook) => [notebook.id, notebook]))
+
+  let notebookCount = 0
+  let cellCount = 0
+
+  const assignCellIds = (cells: OperationResultCell[]): CellWire[] =>
+    cells.map((cell): CellWire => {
+      if ('id' in cell) return cell
+      const id = `mock-cell-${++cellCount}`
+      switch (cell._tag) {
+        case 'markdown_cell':
+          return { ...cell, id }
+        case 'database_cell':
+          return { ...cell, id }
+        case 'log_cell':
+          return { ...cell, id }
+      }
+    })
+
+  return {
+    list: () => [...notebooks.values()],
+    get: (id: string) => notebooks.get(id),
+    create: ({
+      name,
+      description,
+      content,
+    }: {
+      name: string
+      description?: string
+      content: AgentNotebook
+    }) => {
+      const notebook: MockNotebook = {
+        id: `mock-notebook-${++notebookCount}`,
+        name,
+        description,
+        visibility: 'project',
+        updated_at: MOCK_NOTEBOOK_TIMESTAMP,
+        content: { schema_version: content.schema_version, cells: assignCellIds(content.cells) },
+      }
+      notebooks.set(notebook.id, notebook)
+      return notebook
+    },
+    replaceCells: (id: string, cells: OperationResultCell[]) => {
+      const existing = notebooks.get(id)
+      if (!existing) return
+      notebooks.set(id, {
+        ...existing,
+        content: { schema_version: existing.content.schema_version, cells: assignCellIds(cells) },
+      })
+    },
+  }
+}
+
+type MockNotebookStore = ReturnType<typeof createMockNotebookStore>
+
+// All four notebook tools are real, locally-defined ai-SDK tools, so wrap them and
+// override only execute/needsApproval — evals must validate the model's arguments
+// against the exact schemas production uses (agentCellSchema's `.strict()` rejection of
+// agent-authored cell ids, update_notebook's real operations schema, etc).
+function createMockNotebookTools(store: MockNotebookStore) {
+  const { list_notebooks, get_notebook, create_notebook, update_notebook } = getNotebookTools()
+
+  return {
+    list_notebooks: {
+      ...list_notebooks,
+      execute: async (
+        { limit = 20 }: { cursor?: string; limit?: number },
+        _options: ToolCallOptions
+      ) => ({
+        notebooks: store
+          .list()
+          .slice(0, limit)
+          .map((notebook) => ({
+            id: notebook.id,
+            name: notebook.name,
+            description: notebook.description,
+            visibility: notebook.visibility,
+            updated_at: notebook.updated_at,
+            cell_count: notebook.content.cells.length,
+          })),
+        // The in-memory store never paginates: one page holds everything.
+        cursor: undefined,
+      }),
+    },
+    get_notebook: {
+      ...get_notebook,
+      execute: async ({ id }: { id: string }, _options: ToolCallOptions) => {
+        const notebook = store.get(id)
+        if (!notebook) throw new Error(`Notebook ${id} not found.`)
+
+        return {
+          id: notebook.id,
+          name: notebook.name,
+          description: notebook.description,
+          visibility: notebook.visibility,
+          cells: notebook.content.cells,
+        }
+      },
+    },
+    create_notebook: {
+      ...create_notebook,
+      // The eval harness can't answer an approval gate (generate-assistant-response
+      // drops tool parts in 'approval-requested' state when cleaning messages), so the
+      // real needsApproval: true would stall the eval turn — same override as
+      // execute_sql/deploy_edge_function above. Because that gate is gone, this mock
+      // deliberately skips acceptUntrustedSql/acceptUntrustedLogsSql promotion: nothing
+      // here is executed or sent anywhere, cells are stored as plain data in a Map.
+      needsApproval: false,
+      execute: async (
+        {
+          name,
+          description,
+          content,
+        }: {
+          name: string
+          description?: string
+          content: AgentNotebook
+        },
+        _options: ToolCallOptions
+      ) => {
+        const created = store.create({ name, description, content })
+        return { id: created.id, name: created.name }
+      },
+    },
+    update_notebook: {
+      ...update_notebook,
+      // Same reasoning as create_notebook's override above.
+      needsApproval: false,
+      execute: async (
+        { id, operations }: { id: string; operations: NotebookOperation[] },
+        _options: ToolCallOptions
+      ) => {
+        const notebook = store.get(id)
+        if (!notebook) throw new Error(`Notebook ${id} not found.`)
+
+        const result = applyNotebookOperations(notebook.content, operations)
+        if (!result.success) throw new Error(describeNotebookOperationError(result.error))
+
+        store.replaceCells(id, result.notebook.cells)
+        return { id, name: notebook.name }
+      },
+    },
+  }
+}
+
 export type MockToolOverrides = {
   list_tables?: Record<string, typeof MOCK_TABLES_DATA>
 }
@@ -280,6 +507,7 @@ export type MockToolOverrides = {
  */
 export async function getMockTools(overrides: MockToolOverrides | undefined, signal: AbortSignal) {
   const mockedStudioTools = createMockedStudioTools()
+  const notebookStore = createMockNotebookStore()
 
   // Every tool here is a deterministic mock except `search_docs`, which uses the
   // real implementation. We source it from an in-process MCP server directly
@@ -309,5 +537,6 @@ export async function getMockTools(overrides: MockToolOverrides | undefined, sig
     get_advisors: createMockGetAdvisorsTool(),
     query_logs: createMockQueryLogsTool(),
     list_policies: createMockListPoliciesTool(),
+    ...createMockNotebookTools(notebookStore),
   }
 }

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

@@ -5,8 +5,8 @@ import { z } from 'zod'
 import { getContent } from '@/data/content/content-infinite-query'
 import {
   applyNotebookOperations,
+  describeNotebookOperationError,
   notebookOperationsSchema,
-  type NotebookOperationError,
 } from '@/data/content/notebooks/notebook-operations'
 import { getNotebook } from '@/data/content/notebooks/notebook-query'
 import {
@@ -20,17 +20,6 @@ import { createNotebook, updateNotebook } from '@/data/content/notebooks/noteboo
 import { acceptUntrustedLogsSql, untrustedLogSql } from '@/data/logs/safe-analytics-sql'
 import type { Notebooks } from '@/types'
 
-function describeOperationError(error: NotebookOperationError): string {
-  switch (error._tag) {
-    case 'unknown_cell_id':
-      return `No cell with id "${error.cell_id}" exists in this notebook.`
-    case 'conflicting_operations':
-      return `More than one operation targets cell "${error.cell_id}".`
-    case 'empty_result':
-      return 'This update would leave the notebook with no cells.'
-  }
-}
-
 export type NotebookToolsContext = {
   projectRef?: string
   authorization?: string
@@ -191,7 +180,7 @@ export const getNotebookTools = (ctx: NotebookToolsContext = {}) => {
 
         const result = applyNotebookOperations(wireNotebook, operations)
         if (!result.success) {
-          throw new Error(describeOperationError(result.error))
+          throw new Error(describeNotebookOperationError(result.error))
         }
 
         // Same promotion as create_notebook above, inlined here for the same auditability