konnosDashboard
mock-tools.ts330 lines · main
1import assert from 'node:assert'
2import { tool } from 'ai'
3import { z } from 'zod'
4
5import { getStudioTools } from '../tools/studio-tools'
6import { getMcpTools } from '@/lib/ai/tools/mcp-tools'
7
8const listTablesInputSchema = z.object({
9 schemas: z.array(z.string()).describe('The schema names to list.'),
10})
11
12const getAdvisorsInputSchema = z.object({
13 type: z.enum(['security', 'performance']).optional(),
14})
15
16const getLogsInputSchema = z.object({
17 limit: z.number().min(1).max(100).optional(),
18 level: z.enum(['debug', 'info', 'warning', 'error']).optional(),
19 source: z.enum(['postgres', 'auth', 'storage', 'edge_function']).optional(),
20 search: z.string().optional(),
21})
22
23const listPoliciesInputSchema = z.object({
24 schemas: z.array(z.string()).describe('The schema names to get the policies for'),
25})
26
27export const MOCK_TABLES_DATA = [
28 {
29 name: 'user_documents',
30 rls_enabled: false,
31 columns: [
32 { name: 'id', data_type: 'bigint' },
33 { name: 'user_id', data_type: 'uuid' },
34 { name: 'title', data_type: 'text' },
35 ],
36 },
37 {
38 name: 'customers',
39 rls_enabled: true,
40 columns: [
41 { name: 'id', data_type: 'uuid' },
42 { name: 'tenant_id', data_type: 'uuid' },
43 { name: 'email', data_type: 'text' },
44 ],
45 },
46 {
47 name: 'projects',
48 rls_enabled: false,
49 columns: [
50 { name: 'id', data_type: 'uuid' },
51 { name: 'organization_id', data_type: 'uuid' },
52 { name: 'name', data_type: 'text' },
53 ],
54 },
55 {
56 name: 'user_organizations',
57 rls_enabled: true,
58 columns: [
59 { name: 'user_id', data_type: 'uuid' },
60 { name: 'organization_id', data_type: 'uuid' },
61 ],
62 },
63]
64
65const MOCK_EXTENSIONS_DATA = [
66 { name: 'pgcrypto', schema: 'extensions', installed_version: '1.3' },
67 { name: 'uuid-ossp', schema: 'extensions', installed_version: '1.1' },
68 { name: 'pg_cron', schema: 'pg_catalog', installed_version: '1.6.4' },
69]
70
71const MOCK_EDGE_FUNCTIONS_DATA = [
72 { name: 'hello-world', last_deployed_at: '2024-06-10T12:30:00Z' },
73 { name: 'daily-metrics-sync', last_deployed_at: '2024-06-18T08:15:00Z' },
74 { name: 'select-from-table-with-auth-rls', last_deployed_at: '2024-06-19T09:20:00Z' },
75]
76
77const MOCK_ADVISORIES_DATA = [
78 {
79 id: '0016_materialized_view_in_api',
80 level: 'warning',
81 category: 'security',
82 message: 'Materialized views in API schema can bypass RLS. Move them to private schema.',
83 remediationUrl:
84 'https://supabase.com/docs/guides/database/database-advisors?queryGroups=lint&lint=0016_materialized_view_in_api',
85 },
86 {
87 id: '0031_functions_no_rls_guard',
88 level: 'notice',
89 category: 'security',
90 message: 'Function api.health_check should verify auth context before querying tables.',
91 remediationUrl:
92 'https://supabase.com/docs/guides/database/database-advisors?queryGroups=lint&lint=0031_functions_no_rls_guard',
93 },
94 {
95 id: '1012_slow_query',
96 level: 'info',
97 category: 'performance',
98 message:
99 'Query on table edge_function_logs exceeded 3s average execution time over the last hour.',
100 remediationUrl: 'https://supabase.com/docs/guides/platform/performance-advisors#slow-queries',
101 },
102]
103
104const MOCK_LOGS_DATA = [
105 {
106 id: 'log-001',
107 timestamp: '2024-06-20T14:12:00Z',
108 level: 'error',
109 source: 'edge_function' as const,
110 target: 'hello-world',
111 message: "TypeError: fetch failed at await briven.functions.invoke('analytics')",
112 },
113 {
114 id: 'log-002',
115 timestamp: '2024-06-20T14:05:30Z',
116 level: 'warning',
117 source: 'postgres' as const,
118 target: 'connection_pool',
119 message: 'Query timeout exceeded for statement SELECT * FROM public.audit_log_entries',
120 },
121 {
122 id: 'log-003',
123 timestamp: '2024-06-20T13:59:10Z',
124 level: 'info',
125 source: 'edge_function' as const,
126 target: 'daily-metrics-sync',
127 message: 'Invocation completed in 520ms',
128 },
129 {
130 id: 'log-004',
131 timestamp: '2024-06-20T13:50:00Z',
132 level: 'error',
133 source: 'postgres' as const,
134 target: 'trigger:refresh_materialized_views',
135 message: 'permission denied for relation user_documents',
136 },
137 {
138 id: 'log-005',
139 timestamp: '2024-06-20T13:45:00Z',
140 level: 'info',
141 source: 'auth' as const,
142 target: 'email-confirmation',
143 message: 'Sent verification email to alex@example.com',
144 },
145]
146
147function createMockedStudioTools() {
148 const studioTools = getStudioTools()
149
150 return Object.fromEntries(
151 Object.entries(studioTools).map(([name, baseTool]) => {
152 // Always mock execute_sql and deploy_edge_function with needsApproval disabled
153 if (name === 'execute_sql') {
154 return [name, { ...baseTool, needsApproval: false, execute: async () => [] as unknown[] }]
155 }
156 if (name === 'deploy_edge_function') {
157 return [
158 name,
159 { ...baseTool, needsApproval: false, execute: async () => ({ success: true }) },
160 ]
161 }
162 if (typeof baseTool.execute === 'function') {
163 return [name, baseTool]
164 }
165
166 return [
167 name,
168 { ...baseTool, execute: async () => ({ status: 'Tool call mocked successfully.' }) },
169 ]
170 })
171 ) as typeof studioTools
172}
173
174function createMockListTablesTool(overrideData?: Record<string, typeof MOCK_TABLES_DATA>) {
175 return tool({
176 description: 'Lists tables and columns for the provided schemas.',
177 inputSchema: listTablesInputSchema,
178 execute: async ({ schemas }: { schemas: string[] }) => {
179 const effectiveSchemas = schemas?.length ? schemas : ['public']
180 return effectiveSchemas.map((schema) => ({
181 schema,
182 tables: overrideData?.[schema] ?? MOCK_TABLES_DATA,
183 }))
184 },
185 })
186}
187
188function createMockListExtensionsTool() {
189 return tool({
190 description: 'Lists installed database extensions.',
191 inputSchema: z.object({}),
192 execute: async () => {
193 return MOCK_EXTENSIONS_DATA
194 },
195 })
196}
197
198function createMockListEdgeFunctionsTool() {
199 return tool({
200 description: 'Lists available Briven Edge Functions.',
201 inputSchema: z.object({}),
202 execute: async () => {
203 return MOCK_EDGE_FUNCTIONS_DATA
204 },
205 })
206}
207
208function createMockGetAdvisorsTool() {
209 return tool({
210 description: 'Returns advisory notices for the project (mocked).',
211 inputSchema: getAdvisorsInputSchema,
212 execute: async ({ type }: { type?: 'security' | 'performance' }) => {
213 if (type) {
214 return MOCK_ADVISORIES_DATA.filter((advisory) => advisory.category === type)
215 }
216 return MOCK_ADVISORIES_DATA
217 },
218 })
219}
220
221function createMockGetLogsTool() {
222 return tool({
223 description: 'Fetches recent project logs for debugging or health checks (mocked).',
224 inputSchema: getLogsInputSchema,
225 execute: async ({
226 limit = 10,
227 level,
228 source,
229 search,
230 }: {
231 limit?: number
232 level?: 'debug' | 'info' | 'warning' | 'error'
233 source?: 'postgres' | 'auth' | 'storage' | 'edge_function'
234 search?: string
235 }) => {
236 let filtered = MOCK_LOGS_DATA
237
238 if (level) {
239 filtered = filtered.filter((entry) => entry.level === level)
240 }
241
242 if (source) {
243 filtered = filtered.filter((entry) => entry.source === source)
244 }
245
246 if (search) {
247 const needle = search.toLowerCase()
248 filtered = filtered.filter((entry) =>
249 `${entry.message} ${entry.target}`.toLowerCase().includes(needle)
250 )
251 }
252
253 return filtered.slice(0, limit)
254 },
255 })
256}
257
258function createMockListPoliciesTool() {
259 return tool({
260 description: 'Get existing RLS policies for provided schemas.',
261 inputSchema: listPoliciesInputSchema,
262 execute: async ({ schemas }: { schemas: string[] }) => {
263 const effectiveSchemas = schemas?.length ? schemas : ['public']
264 const results = [] as Array<{
265 schema: string
266 table: string
267 policies: Array<{
268 name: string
269 command: 'select' | 'insert' | 'update' | 'delete'
270 using?: string
271 check?: string
272 }>
273 }>
274
275 for (const schema of effectiveSchemas) {
276 if (schema !== 'public') continue
277 results.push(
278 {
279 schema,
280 table: 'customers',
281 policies: [
282 {
283 name: 'customers_tenant_select',
284 command: 'select',
285 using: "(auth.jwt() ->> 'tenant_id')::uuid = tenant_id",
286 },
287 ],
288 },
289 { schema, table: 'user_documents', policies: [] },
290 { schema, table: 'projects', policies: [] }
291 )
292 }
293 return results
294 },
295 })
296}
297
298export type MockToolOverrides = {
299 list_tables?: Record<string, typeof MOCK_TABLES_DATA>
300}
301
302/**
303 * Deterministic mock implementations of MCP/platform tools for evals.
304 * These mirror tool names used in prompts so the model can call them,
305 * but return stable, static data for repeatable tests.
306 *
307 * Note: search_docs uses the real implementation
308 */
309export async function getMockTools(overrides?: MockToolOverrides) {
310 const mockedStudioTools = createMockedStudioTools()
311
312 const { search_docs } = await getMcpTools({
313 accessToken: 'mock-access-token',
314 projectRef: 'mock-project-ref',
315 aiOptInLevel: 'schema_and_log_and_data',
316 })
317
318 assert(search_docs, 'search_docs tool not available from MCP server')
319
320 return {
321 ...mockedStudioTools,
322 search_docs,
323 list_tables: createMockListTablesTool(overrides?.list_tables),
324 list_extensions: createMockListExtensionsTool(),
325 list_edge_functions: createMockListEdgeFunctionsTool(),
326 get_advisors: createMockGetAdvisorsTool(),
327 get_logs: createMockGetLogsTool(),
328 list_policies: createMockListPoliciesTool(),
329 }
330}