import assert from 'node:assert' import { tool, type ToolExecutionOptions } 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 { createSearchDocsTool } from '@/lib/ai/tools/search-docs-tool' const listTablesInputSchema = z.object({ schemas: z.array(z.string()).describe('The schema names to list.'), }) const getAdvisorsInputSchema = z.object({ type: z.enum(['security', 'performance']).optional(), }) const queryLogsInputSchema = z.object({ sql: z.string().min(1), iso_timestamp_start: z.string().optional(), iso_timestamp_end: z.string().optional(), }) const listPoliciesInputSchema = z.object({ schemas: z.array(z.string()).describe('The schema names to get the policies for'), }) export const MOCK_TABLES_DATA = [ { name: 'user_documents', rls_enabled: false, columns: [ { name: 'id', data_type: 'bigint' }, { name: 'user_id', data_type: 'uuid' }, { name: 'title', data_type: 'text' }, ], }, { name: 'customers', rls_enabled: true, columns: [ { name: 'id', data_type: 'uuid' }, { name: 'tenant_id', data_type: 'uuid' }, { name: 'email', data_type: 'text' }, ], }, { name: 'projects', rls_enabled: false, columns: [ { name: 'id', data_type: 'uuid' }, { name: 'organization_id', data_type: 'uuid' }, { name: 'name', data_type: 'text' }, ], }, { name: 'user_organizations', rls_enabled: true, columns: [ { name: 'user_id', data_type: 'uuid' }, { name: 'organization_id', data_type: 'uuid' }, ], }, ] const MOCK_EXTENSIONS_DATA = [ { name: 'pgcrypto', schema: 'extensions', installed_version: '1.3' }, { name: 'uuid-ossp', schema: 'extensions', installed_version: '1.1' }, { name: 'pg_cron', schema: 'pg_catalog', installed_version: '1.6.4' }, ] const MOCK_EDGE_FUNCTIONS_DATA = [ { name: 'hello-world', last_deployed_at: '2024-06-10T12:30:00Z' }, { name: 'daily-metrics-sync', last_deployed_at: '2024-06-18T08:15:00Z' }, { name: 'select-from-table-with-auth-rls', last_deployed_at: '2024-06-19T09:20:00Z' }, ] const MOCK_ADVISORIES_DATA = [ { id: '0016_materialized_view_in_api', level: 'warning', category: 'security', message: 'Materialized views in API schema can bypass RLS. Move them to private schema.', remediationUrl: 'https://supabase.com/docs/guides/observability/advisors?queryGroups=lint&lint=0016_materialized_view_in_api', }, { id: '0031_functions_no_rls_guard', level: 'notice', category: 'security', message: 'Function api.health_check should verify auth context before querying tables.', remediationUrl: 'https://supabase.com/docs/guides/observability/advisors?queryGroups=lint&lint=0031_functions_no_rls_guard', }, { id: '1012_slow_query', level: 'info', category: 'performance', message: 'Query on table edge_function_logs exceeded 3s average execution time over the last hour.', remediationUrl: 'https://supabase.com/docs/guides/platform/performance-advisors#slow-queries', }, ] const MOCK_LOGS_DATA = [ { id: 'log-001', timestamp: '2024-06-20T14:12:00Z', level: 'error', source: 'edge_function' as const, target: 'hello-world', message: "TypeError: fetch failed at await supabase.functions.invoke('analytics')", }, { id: 'log-002', timestamp: '2024-06-20T14:05:30Z', level: 'warning', source: 'postgres' as const, target: 'connection_pool', message: 'Query timeout exceeded for statement SELECT * FROM public.audit_log_entries', }, { id: 'log-003', timestamp: '2024-06-20T13:59:10Z', level: 'info', source: 'edge_function' as const, target: 'daily-metrics-sync', message: 'Invocation completed in 520ms', }, { id: 'log-004', timestamp: '2024-06-20T13:50:00Z', level: 'error', source: 'postgres' as const, target: 'trigger:refresh_materialized_views', message: 'permission denied for relation user_documents', }, { id: 'log-005', timestamp: '2024-06-20T13:45:00Z', level: 'info', source: 'auth' as const, target: 'email-confirmation', message: 'Sent verification email to alex@example.com', }, ] 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_series: ['signups'], cumulative: false, type: 'line', scale: 'log', show_labels: 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() return Object.fromEntries( Object.entries(studioTools).map(([name, baseTool]) => { // Always mock execute_sql and deploy_edge_function with needsApproval disabled if (name === 'execute_sql') { return [name, { ...baseTool, needsApproval: false, execute: async () => [] as unknown[] }] } if (name === 'deploy_edge_function') { return [ name, { ...baseTool, needsApproval: false, execute: async () => ({ success: true }) }, ] } if (typeof baseTool.execute === 'function') { return [name, baseTool] } return [ name, { ...baseTool, execute: async () => ({ status: 'Tool call mocked successfully.' }) }, ] }) ) as typeof studioTools } function createMockListTablesTool(overrideData?: Record) { return tool({ description: 'Lists tables and columns for the provided schemas.', inputSchema: listTablesInputSchema, execute: async ({ schemas }: { schemas: string[] }) => { const effectiveSchemas = schemas?.length ? schemas : ['public'] return effectiveSchemas.map((schema) => ({ schema, tables: overrideData?.[schema] ?? MOCK_TABLES_DATA, })) }, }) } function createMockListExtensionsTool() { return tool({ description: 'Lists installed database extensions.', inputSchema: z.object({}), execute: async () => { return MOCK_EXTENSIONS_DATA }, }) } function createMockListEdgeFunctionsTool() { return tool({ description: 'Lists available Supabase Edge Functions.', inputSchema: z.object({}), execute: async () => { return MOCK_EDGE_FUNCTIONS_DATA }, }) } function createMockGetAdvisorsTool() { return tool({ description: 'Returns advisory notices for the project (mocked).', inputSchema: getAdvisorsInputSchema, execute: async ({ type }: { type?: 'security' | 'performance' }) => { if (type) { return MOCK_ADVISORIES_DATA.filter((advisory) => advisory.category === type) } return MOCK_ADVISORIES_DATA }, }) } function createMockQueryLogsTool() { return tool({ description: 'Runs a read-only SQL query against recent project logs for debugging or health checks (mocked).', inputSchema: queryLogsInputSchema, // Deterministic mock: returns static log data regardless of the SQL passed. execute: async () => MOCK_LOGS_DATA, }) } function createMockListPoliciesTool() { return tool({ description: 'Get existing RLS policies for provided schemas.', inputSchema: listPoliciesInputSchema, execute: async ({ schemas }: { schemas: string[] }) => { const effectiveSchemas = schemas?.length ? schemas : ['public'] const results = [] as Array<{ schema: string table: string policies: Array<{ name: string command: 'select' | 'insert' | 'update' | 'delete' using?: string check?: string }> }> for (const schema of effectiveSchemas) { if (schema !== 'public') continue results.push( { schema, table: 'customers', policies: [ { name: 'customers_tenant_select', command: 'select', using: "(auth.jwt() ->> 'tenant_id')::uuid = tenant_id", }, ], }, { schema, table: 'user_documents', policies: [] }, { schema, table: 'projects', policies: [] } ) } return results }, }) } 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), delete: (id: string) => notebooks.delete(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 const MOCK_DATABASES_DATA = [ { identifier: 'mock-project-ref', is_primary: true, region: 'us-east-1', status: 'ACTIVE_HEALTHY', }, { identifier: 'mock-project-ref-replica-1', is_primary: false, region: 'us-west-1', status: 'ACTIVE_HEALTHY', }, ] // All 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_databases, list_notebooks, get_notebook, run_notebook, create_notebook, update_notebook, delete_notebook, } = getNotebookTools({ aiOptInLevel: 'schema_and_log_and_data' }) return { list_databases: { ...list_databases, execute: async (_args: object, _options: ToolExecutionOptions) => ({ databases: MOCK_DATABASES_DATA, }), }, list_notebooks: { ...list_notebooks, execute: async ( { limit = 20 }: { cursor?: string; limit?: number }, _options: ToolExecutionOptions ) => ({ 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: ToolExecutionOptions) => { 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, updated_at: notebook.updated_at, cells: notebook.content.cells, } }, }, run_notebook: { ...run_notebook, // The eval harness cannot answer approval gates. Nothing executes here; return a // deterministic empty result for each query cell in notebook order. needsApproval: false, execute: async ( { id }: { id: string; expected_updated_at: string }, _options: ToolExecutionOptions ) => { const notebook = store.get(id) if (!notebook) throw new Error(`Notebook ${id} not found.`) return { id, name: notebook.name, updated_at: notebook.updated_at, cells: notebook.content.cells.flatMap((cell) => cell._tag === 'markdown_cell' ? [] : [ { cell_id: cell._id, title: cell.title?.trim() || 'Untitled query', source: cell._tag === 'log_cell' ? ('logs' as const) : ('database' as const), status: 'success' as const, rows: [], }, ] ), } }, }, 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: ToolExecutionOptions ) => { 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, // expected_updated_at is validated by the real inputSchema (spread above) but not // checked here: the in-memory store has no concurrent writers for the eval harness // to race against. execute: async ( { id, operations, }: { id: string; expected_updated_at: string; operations: NotebookOperation[] }, _options: ToolExecutionOptions ) => { 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 } }, }, delete_notebook: { ...delete_notebook, // Same reasoning as create_notebook's override above. needsApproval: false, execute: async ({ id }: { id: string }, _options: ToolExecutionOptions) => { const notebook = store.get(id) if (!notebook) throw new Error(`Notebook ${id} not found.`) store.delete(id) return { id, name: notebook.name } }, }, } } export type MockToolOverrides = { list_tables?: Record } /** * Deterministic mock implementations of MCP/platform tools for evals. * These mirror tool names used in prompts so the model can call them, * but return stable, static data for repeatable tests. * * Note: search_docs uses the real implementation. */ export async function getMockTools(overrides: MockToolOverrides | undefined) { const mockedStudioTools = createMockedStudioTools() const notebookStore = createMockNotebookStore() const search_docs = await createSearchDocsTool() const tools = { ...mockedStudioTools, search_docs, list_tables: createMockListTablesTool(overrides?.list_tables), list_extensions: createMockListExtensionsTool(), list_edge_functions: createMockListEdgeFunctionsTool(), get_advisors: createMockGetAdvisorsTool(), query_logs: createMockQueryLogsTool(), list_policies: createMockListPoliciesTool(), ...createMockNotebookTools(notebookStore), } assert(tools.search_docs, 'search_docs tool is missing from the eval harness') return tools }