diff --git a/apps/studio/lib/ai/evals/generate-task.ts b/apps/studio/lib/ai/evals/generate-task.ts new file mode 100644 index 00000000000..3f0c65d6bd8 --- /dev/null +++ b/apps/studio/lib/ai/evals/generate-task.ts @@ -0,0 +1,86 @@ +import { generateText, stepCountIs, type ModelMessage, type ToolSet } from 'ai' +import { source } from 'common-tags' + +import { AiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi' +import { getModel } from 'lib/ai/model' +import { getTools } from 'lib/ai/tools' +import { getRenderingTools } from 'lib/ai/tools/rendering-tools' +import { getFallbackTools } from 'lib/ai/tools/fallback-tools' +import { getEvalMockTools } from './mock-tools' +import { + CHAT_PROMPT, + EDGE_FUNCTION_PROMPT, + GENERAL_PROMPT, + PG_BEST_PRACTICES, + RLS_PROMPT, + SECURITY_PROMPT, +} from 'lib/ai/prompts' + +type GenerateTaskOptions = { + projectRef?: string + chatName?: string + isLimited?: boolean +} + +/** + * Mirrors the generate-v4 flow for task generation, but returns a simple text result + * using generateText instead of streamText. + */ +export async function generateTask(input: string, opts: GenerateTaskOptions = {}) { + const { projectRef = 'eval-project', chatName = 'RLS Eval', isLimited = true } = opts + + const { + model, + error: modelError, + promptProviderOptions, + providerOptions, + } = await getModel({ + provider: 'openai', + model: 'gpt-5-mini', + routingKey: projectRef, + isLimited, + }) + + if (modelError || !model) { + throw modelError ?? new Error('Model unavailable') + } + + // Build system prompt like generate-v4 + const system = source` + ${GENERAL_PROMPT} + ${CHAT_PROMPT} + ${PG_BEST_PRACTICES} + ${RLS_PROMPT} + ${EDGE_FUNCTION_PROMPT} + ${SECURITY_PROMPT} + ` + + const coreMessages: ModelMessage[] = [ + { + role: 'system', + content: system, + ...(promptProviderOptions && { providerOptions: promptProviderOptions }), + }, + { + role: 'assistant', + content: `The user's current project is ${projectRef}. Their available schemas are: You don't have access to any schemas. The current chat name is: ${chatName}`, + }, + { role: 'user', content: input }, + ] + + // Deterministic tools for evals: always include rendering tools + mocked MCP/platform tools + let tools: ToolSet = { + ...getRenderingTools(), + ...getEvalMockTools(), + } + + const result = await generateText({ + model, + messages: coreMessages, + stopWhen: stepCountIs(5), + ...(providerOptions && { providerOptions }), + tools, + }) + + return result.text +} diff --git a/apps/studio/lib/ai/evals/mock-tools.ts b/apps/studio/lib/ai/evals/mock-tools.ts new file mode 100644 index 00000000000..7ba73aa5b15 --- /dev/null +++ b/apps/studio/lib/ai/evals/mock-tools.ts @@ -0,0 +1,107 @@ +import { tool, type ToolSet } from 'ai' +import { z } from 'zod' + +/** + * 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. + */ +export function getEvalMockTools(): ToolSet { + return { + list_tables: tool({ + description: 'Lists tables and columns for the provided schemas.', + inputSchema: z.object({ + schemas: z.array(z.string()).describe('The schema names to list.'), + }), + execute: async ({ schemas }: { schemas: string[] }) => { + const effectiveSchemas = schemas?.length ? schemas : ['public'] + // Deterministic, minimal schema metadata suitable for RLS generation + const result = effectiveSchemas.map((schema) => ({ + schema, + tables: [ + { + name: 'user_documents', + columns: [ + { name: 'id', data_type: 'bigint' }, + { name: 'user_id', data_type: 'uuid' }, + { name: 'title', data_type: 'text' }, + ], + }, + { + name: 'customers', + columns: [ + { name: 'id', data_type: 'uuid' }, + { name: 'tenant_id', data_type: 'uuid' }, + { name: 'email', data_type: 'text' }, + ], + }, + { + name: 'projects', + columns: [ + { name: 'id', data_type: 'uuid' }, + { name: 'organization_id', data_type: 'uuid' }, + { name: 'name', data_type: 'text' }, + ], + }, + ], + })) + return result + }, + }), + + list_extensions: tool({ + description: 'Lists installed database extensions.', + inputSchema: z.object({}), + execute: async () => { + return [ + { name: 'pgcrypto', schema: 'extensions', installed_version: '1.3' }, + { name: 'uuid-ossp', schema: 'extensions', installed_version: '1.1' }, + ] + }, + }), + + list_edge_functions: tool({ + description: 'Lists available Supabase Edge Functions.', + inputSchema: z.object({}), + execute: async () => { + return [{ name: 'hello-world' }, { name: 'select-from-table-with-auth-rls' }] + }, + }), + + get_advisors: tool({ + description: 'Returns advisory notices for the project (mocked).', + inputSchema: z.object({ + type: z.enum(['security', 'performance']).optional(), + }), + execute: async ({ type }: { type?: 'security' | 'performance' }) => { + if (type === 'security') { + return [ + { + id: '0016_materialized_view_in_api', + level: 'warning', + message: + 'Materialized views in API schema can bypass RLS. Move them to private schema.', + remediationUrl: + 'https://supabase.com/docs/guides/database/database-advisors?queryGroups=lint&lint=0016_materialized_view_in_api', + }, + ] + } + return [] + }, + }), + + list_policies: tool({ + description: 'Get existing RLS policies for provided schemas.', + inputSchema: z.object({ + schemas: z.array(z.string()).describe('The schema names to get the policies for'), + }), + execute: async ({ schemas }: { schemas: string[] }) => { + // Keep simple and deterministic: no existing policies + const effectiveSchemas = schemas?.length ? schemas : ['public'] + return effectiveSchemas + .map((s) => `No existing policies in schema ${s}. You can create new policies as needed.`) + .join('\n') + }, + }), + } +} diff --git a/apps/studio/lib/ai/evals/rls.eval.ts b/apps/studio/lib/ai/evals/rls.eval.ts index e69de29bb2d..f79d9bef831 100644 --- a/apps/studio/lib/ai/evals/rls.eval.ts +++ b/apps/studio/lib/ai/evals/rls.eval.ts @@ -0,0 +1,43 @@ +import { Eval } from 'braintrust' +import { Levenshtein } from 'autoevals' +import { generateTask } from './generate-task' + +// Evaluate how well the generate-v4-like function creates RLS policies +Eval('RLS Policy Generation', { + data: () => { + return [ + // Simple single-table, user-ownership pattern + { + input: + 'Create RLS policies for table public.user_documents with column user_id so users can read/write only their own rows. Include select, insert, update, delete.', + expected: 'CREATE POLICY', // We check fuzzy string similarity; exact SQL compared by Levenshtein + }, + // Multi-tenant via jwt claim pattern + { + input: + "For table public.customers with tenant_id, add SELECT policy restricting rows to tenant_id from auth.jwt() ->> 'tenant_id'.", + expected: 'CREATE POLICY', + }, + // Org membership via join-like pattern + { + input: + 'For table public.projects with organization_id, allow SELECT for members listed in public.user_organizations(user_id, organization_id) matched by auth.uid().', + expected: 'CREATE POLICY', + }, + // Storage example + { + input: + "For table storage.objects, allow SELECT/INSERT only for bucket 'user-uploads' where first folder equals auth.uid().", + expected: 'CREATE POLICY', + }, + ] + }, + task: async (input: string) => { + // Use default options suitable for eval environment (no schema access) + const output = await generateTask(input, { + isLimited: false, + }) + return output + }, + scores: [Levenshtein], +})