diff --git a/apps/studio/data/content/notebooks/notebook-operations.ts b/apps/studio/data/content/notebooks/notebook-operations.ts index 8b9ba411ddc..8ff279ac980 100644 --- a/apps/studio/data/content/notebooks/notebook-operations.ts +++ b/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': diff --git a/apps/studio/lib/ai/tools/mock-tools.test.ts b/apps/studio/lib/ai/tools/mock-tools.test.ts index 9e86efbb9c4..a9b808b703f 100644 --- a/apps/studio/lib/ai/tools/mock-tools.test.ts +++ b/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', + ]) + }) + }) }) diff --git a/apps/studio/lib/ai/tools/mock-tools.ts b/apps/studio/lib/ai/tools/mock-tools.ts index 049b0b5fecb..2ceb1b452b2 100644 --- a/apps/studio/lib/ai/tools/mock-tools.ts +++ b/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 + +// 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 } @@ -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), } } diff --git a/apps/studio/lib/ai/tools/notebook-tools.ts b/apps/studio/lib/ai/tools/notebook-tools.ts index 64ae842fd24..222ec181175 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.ts +++ b/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